arize-phoenix 0.0.32__py3-none-any.whl → 0.0.33__py3-none-any.whl

phoenix/server/thread_server.py

@@ -13,10 +13,12 @@ class ThreadServer(Server):
     def __init__(
         self,
         app: Starlette,
+        host: str,
         port: int,
     ) -> None:
         config = Config(
             app=app,
+            host=host,
             port=port,
             # TODO: save logs to file
             log_level=logging.ERROR,
@@ -32,9 +34,15 @@ class ThreadServer(Server):
         thread.start()
         time_limit = time() + 5  # 5 seconds
         try:
-            while not self.started and thread.is_alive() and time() < time_limit:
+            while (
+                time() < time_limit
+                and thread.is_alive()
+                and not self.should_exit
+                and not self.started
+            ):
                 sleep(1e-3)
-            if time() > time_limit:
+            if time() >= time_limit and not self.started:
+                self.should_exit = True
                 raise RuntimeError("server took too long to start")
             yield thread
         finally:

phoenix/services.py

@@ -4,11 +4,12 @@ import signal
 import subprocess
 import sys
 from pathlib import Path
-from typing import List, Optional
+from time import sleep, time
+from typing import Callable, List, Optional
 
 import psutil
 
-import phoenix.config as config
+from phoenix.config import SERVER_DIR, get_pids_path, get_running_pid
 
 logger = logging.getLogger(__name__)
 
@@ -22,6 +23,12 @@ class Service:
 
     def __init__(self) -> None:
         self.child = self.start()
+        self._wait_until(
+            lambda: get_running_pid() is not None,
+            # Not sure why, but the process can take a very long time
+            # to get going, e.g. 15+ seconds in Colab.
+            up_to_seconds=60,
+        )
 
     @property
     def command(self) -> List[str]:
@@ -30,7 +37,7 @@ class Service:
     def start(self) -> psutil.Popen:
         """Starts the service."""
 
-        if len(os.listdir(config.get_pids_path())) > 0:
+        if get_running_pid():
             # Currently, only one instance of Phoenix can be running at any given time.
             # Support for multiple concurrently running instances may be supported in the future.
             logger.warning(
@@ -48,21 +55,18 @@ class Service:
             text=True,
             env={**os.environ},
         )
-        # TODO: convert to async with timeout because this can block forever
-        # if there's nothing to read. This is also brittle because it relies
-        # on a specific line of print output by a third party module (uvicorn).
-        for line in iter(process.stdout.readline, b""):
-            if "Uvicorn running on" in str(line):
-                break
         return process
 
     @property
     def active(self) -> bool:
-        return self.child.is_running()
+        # Not sure why, but the process can remain in a zombie state
+        # indefinitely, e.g. in Colab.
+        return self.child.is_running() and self.child.status() != psutil.STATUS_ZOMBIE
 
     def stop(self) -> None:
         """Stops the service."""
         self.child.terminate()
+        self._wait_until(lambda: get_running_pid() is None)
 
     @staticmethod
     def stop_any() -> None:
@@ -70,39 +74,54 @@ class Service:
         within the current session or if it is being run in a separate process on the
         same host machine. In either case, the instance will be forcibly stopped.
         """
-        pids_path = config.get_pids_path()
-        for filename in os.listdir(pids_path):
+        for file in get_pids_path().iterdir():
+            if not file.name.isnumeric():
+                continue
             try:
-                os.kill(int(filename), signal.SIGKILL)
+                os.kill(int(file.name), signal.SIGKILL)
             except ProcessLookupError:
                 pass
-            filename_path = os.path.join(pids_path, filename)
-            os.unlink(filename_path)
+            file.unlink(missing_ok=True)
+
+    def _wait_until(
+        self,
+        predicate: Callable[[], bool],
+        up_to_seconds: float = 5,
+        sleep_seconds: float = 1e-3,
+    ) -> None:
+        time_limit = time() + up_to_seconds
+        while not predicate() and time() < time_limit and self.active:
+            sleep(sleep_seconds)
 
 
 class AppService(Service):
     """Service that controls the phoenix application."""
 
-    working_dir = config.SERVER_DIR
+    working_dir = SERVER_DIR
 
     # Internal references to the name / directory of the dataset(s)
     __primary_dataset_name: str
     __reference_dataset_name: Optional[str]
     __corpus_dataset_name: Optional[str]
+    __trace_dataset_name: Optional[str]
 
     def __init__(
         self,
         export_path: Path,
+        host: str,
         port: int,
         primary_dataset_name: str,
         reference_dataset_name: Optional[str],
         corpus_dataset_name: Optional[str],
+        trace_dataset_name: Optional[str],
     ):
         self.export_path = export_path
+        self.host = host
         self.port = port
         self.__primary_dataset_name = primary_dataset_name
         self.__reference_dataset_name = reference_dataset_name
         self.__corpus_dataset_name = corpus_dataset_name
+        self.__trace_dataset_name = trace_dataset_name
         super().__init__()
 
     @property
@@ -112,6 +131,8 @@ class AppService(Service):
             "main.py",
             "--export_path",
             str(self.export_path),
+            "--host",
+            str(self.host),
             "--port",
             str(self.port),
             "datasets",
@@ -122,5 +143,7 @@ class AppService(Service):
             command.extend(["--reference", str(self.__reference_dataset_name)])
         if self.__corpus_dataset_name is not None:
             command.extend(["--corpus", str(self.__corpus_dataset_name)])
+        if self.__trace_dataset_name is not None:
+            command.extend(["--trace", str(self.__trace_dataset_name)])
         logger.info(f"command: {command}")
         return command

phoenix/session/session.py

@@ -1,19 +1,24 @@
+import json
 import logging
 from abc import ABC, abstractmethod
 from collections import UserList
+from datetime import datetime
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Iterable, List, Optional, Set
 
 import pandas as pd
-from portpicker import pick_unused_port
 
-from phoenix.config import PORT, get_exported_files
+from phoenix.config import get_env_host, get_env_port, get_exported_files
 from phoenix.core.model_schema_adapter import create_model_from_datasets
-from phoenix.datasets.dataset import Dataset
+from phoenix.core.traces import Traces
+from phoenix.datasets.dataset import EMPTY_DATASET, Dataset
 from phoenix.server.app import create_app
 from phoenix.server.thread_server import ThreadServer
 from phoenix.services import AppService
+from phoenix.trace.filter import SpanFilter
+from phoenix.trace.span_json_encoder import span_to_json
+from phoenix.trace.trace_dataset import TraceDataset
 
 try:
     from IPython.display import IFrame  # type: ignore
@@ -52,6 +57,9 @@ class ExportedData(_BaseList):
 class Session(ABC):
     """Session that maintains a 1-1 shared state with the Phoenix App."""
 
+    trace_dataset: Optional[TraceDataset]
+    traces: Optional[Traces]
+
     def __dir__(self) -> List[str]:
         return ["exports", "view", "url"]
 
@@ -60,11 +68,14 @@ class Session(ABC):
         primary_dataset: Dataset,
         reference_dataset: Optional[Dataset] = None,
         corpus_dataset: Optional[Dataset] = None,
-        port: int = PORT,
+        trace_dataset: Optional[TraceDataset] = None,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
     ):
         self.primary_dataset = primary_dataset
         self.reference_dataset = reference_dataset
         self.corpus_dataset = corpus_dataset
+        self.trace_dataset = trace_dataset
         self.model = create_model_from_datasets(
             primary_dataset,
             reference_dataset,
@@ -78,7 +89,13 @@ class Session(ABC):
             else None
         )
 
-        self.port = port
+        self.traces = Traces()
+        if trace_dataset:
+            for span in trace_dataset.to_spans():
+                self.traces.put(span)
+
+        self.host = host or get_env_host()
+        self.port = port or get_env_port()
         self.temp_dir = TemporaryDirectory()
         self.export_path = Path(self.temp_dir.name) / "exports"
         self.export_path.mkdir(parents=True, exist_ok=True)
@@ -122,7 +139,29 @@ class Session(ABC):
     @property
     def url(self) -> str:
         """Returns the url for the phoenix app"""
-        return _get_url(self.port, self.is_colab)
+        return _get_url(self.host, self.port, self.is_colab)
+
+    def get_span_dataframe(
+        self,
+        filter_condition: Optional[str] = None,
+        *,
+        start_time: Optional[datetime] = None,
+        stop_time: Optional[datetime] = None,
+        root_spans_only: Optional[bool] = None,
+    ) -> Optional[pd.DataFrame]:
+        if (traces := self.traces) is None:
+            return None
+        predicate = SpanFilter(filter_condition) if filter_condition else None
+        spans = traces.get_spans(
+            start_time=start_time,
+            stop_time=stop_time,
+            root_spans_only=root_spans_only,
+        )
+        if predicate:
+            spans = filter(predicate, spans)
+        if not (data := list(map(json.loads, map(span_to_json, spans)))):
+            return None
+        return pd.json_normalize(data).set_index("context.span_id", drop=False)
 
 
 _session: Optional[Session] = None
@@ -134,22 +173,29 @@ class ProcessSession(Session):
         primary_dataset: Dataset,
         reference_dataset: Optional[Dataset] = None,
         corpus_dataset: Optional[Dataset] = None,
+        trace_dataset: Optional[TraceDataset] = None,
+        host: Optional[str] = None,
         port: Optional[int] = None,
     ) -> None:
         super().__init__(
             primary_dataset=primary_dataset,
             reference_dataset=reference_dataset,
             corpus_dataset=corpus_dataset,
-            port=port or PORT,
+            trace_dataset=trace_dataset,
+            host=host,
+            port=port,
         )
         primary_dataset.to_disc()
         if isinstance(reference_dataset, Dataset):
             reference_dataset.to_disc()
         if isinstance(corpus_dataset, Dataset):
             corpus_dataset.to_disc()
+        if isinstance(trace_dataset, TraceDataset):
+            trace_dataset.to_disc()
         # Initialize an app service that keeps the server running
         self.app_service = AppService(
             self.export_path,
+            self.host,
             self.port,
             self.primary_dataset.name,
             reference_dataset_name=(
@@ -158,6 +204,9 @@ class ProcessSession(Session):
             corpus_dataset_name=(
                 self.corpus_dataset.name if self.corpus_dataset is not None else None
             ),
+            trace_dataset_name=(
+                self.trace_dataset.name if self.trace_dataset is not None else None
+            ),
         )
 
     @property
@@ -175,22 +224,28 @@ class ThreadSession(Session):
         primary_dataset: Dataset,
         reference_dataset: Optional[Dataset] = None,
         corpus_dataset: Optional[Dataset] = None,
+        trace_dataset: Optional[TraceDataset] = None,
+        host: Optional[str] = None,
         port: Optional[int] = None,
     ):
         super().__init__(
             primary_dataset=primary_dataset,
             reference_dataset=reference_dataset,
             corpus_dataset=corpus_dataset,
-            port=port or pick_unused_port(),
+            trace_dataset=trace_dataset,
+            host=host,
+            port=port,
         )
         # Initialize an app service that keeps the server running
         self.app = create_app(
             export_path=self.export_path,
             model=self.model,
             corpus=self.corpus,
+            traces=self.traces,
         )
         self.server = ThreadServer(
             app=self.app,
+            host=self.host,
             port=self.port,
         ).run_in_thread()
         # start the server
@@ -206,11 +261,13 @@ class ThreadSession(Session):
 
 
 def launch_app(
-    primary: Dataset,
+    primary: Optional[Dataset] = None,
     reference: Optional[Dataset] = None,
     corpus: Optional[Dataset] = None,
+    trace: Optional[TraceDataset] = None,
+    host: Optional[str] = None,
     port: Optional[int] = None,
-    run_in_thread: Optional[bool] = True,
+    run_in_thread: bool = True,
 ) -> Optional[Session]:
     """
     Launches the phoenix application and returns a session to interact with.
@@ -224,8 +281,14 @@ def launch_app(
         If not provided, drift analysis will not be available.
     corpus : Dataset, optional
         The dataset containing corpus for LLM context retrieval.
+    trace: TraceDataset, optional
+        **Experimental** The trace dataset containing the trace data.
+    host: str, optional
+        The host on which the server runs. It can also be set using environment
+        variable `PHOENIX_HOST`, otherwise it defaults to `127.0.0.1`.
     port: int, optional
-        The port on which the server listens.
+        The port on which the server listens. It can also be set using environment
+        variable `PHOENIX_PORT`, otherwise it defaults to 6060.
     run_in_thread: bool, optional, default=True
         Whether the server should run in a Thread or Process.
 
@@ -243,23 +306,32 @@ def launch_app(
     """
     global _session
 
+    # Stopgap solution to allow the app to run without a primary dataset
+    if primary is None:
+        # Dummy dataset
+        # TODO: pass through the lack of a primary dataset to the app
+        primary = EMPTY_DATASET
+
+    if _session is not None and _session.active:
+        logger.warning(
+            "Existing running Phoenix instance detected! Shutting "
+            "it down and starting a new instance..."
+        )
+        _session.end()
+
     if run_in_thread:
-        if _session is not None and _session.active:
-            logger.warning(
-                "Existing running Phoenix instance detected! Shutting "
-                "it down and starting a new instance..."
-            )
-            _session.end()
-        _session = ThreadSession(primary, reference, corpus, port=port)
+        _session = ThreadSession(primary, reference, corpus, trace, host=host, port=port)
         # TODO: catch exceptions from thread
-        if not _session.active:
-            logger.error(
-                "💥 Phoenix failed to start. Please try again or file an issue "
-                "with us at https://github.com/Arize-ai/phoenix"
-            )
-            return None
     else:
-        _session = ProcessSession(primary, reference, port=port)
+        _session = ProcessSession(primary, reference, corpus, trace, host=host, port=port)
+
+    if not _session.active:
+        logger.error(
+            f"💥 Phoenix failed to start. Please try again (making sure that "
+            f"port {port} is not occupied by another process) or file an issue "
+            f"with us at https://github.com/Arize-ai/phoenix"
+        )
+        return None
 
     print(f"🌍 To view the Phoenix app in your browser, visit {_session.url}")
     print("📺 To view the Phoenix app in a notebook, run `px.active_session().view()`")
@@ -288,14 +360,14 @@ def close_app() -> None:
     logger.info("Session closed")
 
 
-def _get_url(port: int, is_colab: bool) -> str:
+def _get_url(host: str, port: int, is_colab: bool) -> str:
     """Determines the IFrame URL based on whether this is in a Colab or in a local notebook"""
     if is_colab:
         from google.colab.output import eval_js  # type: ignore
 
         return str(eval_js(f"google.colab.kernel.proxyPort({port}, {{'cache': true}})"))
 
-    return f"http://localhost:{port}/"
+    return f"http://{host}:{port}/"
 
 
 def _is_colab() -> bool:

phoenix/trace/exporter.py

@@ -0,0 +1,71 @@
+import gzip
+import logging
+import weakref
+from queue import SimpleQueue
+from threading import Thread
+from types import MethodType
+from typing import Optional
+
+from requests import Session
+
+from phoenix.config import get_env_host, get_env_port
+from phoenix.trace.schemas import Span
+from phoenix.trace.v1 import encode
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+
+
+class NoOpExporter:
+    def export(self, span: Span) -> None:
+        pass
+
+
+class HttpExporter:
+    def __init__(
+        self,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+    ) -> None:
+        self._host = host or get_env_host()
+        self._port = port or get_env_port()
+        self._url = f"http://{self._host}:{self._port}/v1/spans"
+        self._session = Session()
+        weakref.finalize(self, self._session.close)
+        self._session.headers.update(
+            {
+                "content-type": "application/x-protobuf",
+                "content-encoding": "gzip",
+            }
+        )
+        self._queue: "SimpleQueue[Optional[Span]]" = SimpleQueue()
+        # Putting `None` as the sentinel value for queue termination.
+        weakref.finalize(self, self._queue.put, None)
+        self._start_consumer()
+
+    def export(self, span: Span) -> None:
+        self._queue.put(span)
+
+    def _start_consumer(self) -> None:
+        Thread(
+            target=MethodType(
+                self.__class__._consume_spans,
+                weakref.proxy(self),
+            ),
+            daemon=True,
+        ).start()
+
+    def _consume_spans(self) -> None:
+        while True:
+            if not (span := self._queue.get()):
+                return
+            self._send(span)
+
+    def _send(self, span: Span) -> None:
+        pb_span = encode(span)
+        serialized = pb_span.SerializeToString()
+        data = gzip.compress(serialized)
+        try:
+            self._session.post(self._url, data=data)
+        except Exception as e:
+            logger.exception(e)

phoenix/trace/filter.py

@@ -0,0 +1,181 @@
+import ast
+from typing import Any, Iterator, Mapping, Tuple, cast
+
+from phoenix.trace import semantic_conventions
+from phoenix.trace.schemas import Span
+
+
+class SpanFilter:
+    def __init__(self, condition: str) -> None:
+        self._root = ast.parse(condition, mode="eval")
+        _validate_expression(self._root, condition)
+        self._translated = _Translator(condition).visit(self._root)
+        ast.fix_missing_locations(self._translated)
+        self._compiled = compile(self._translated, filename="", mode="eval")
+
+    def __call__(self, span: Span) -> bool:
+        return cast(bool, eval(self._compiled, {}, {"span": span, "_MISSING": _Missing()}))
+
+
+def _replace_none_with_missing(
+    value: ast.expr,
+    as_str: bool = False,
+) -> ast.IfExp:
+    """
+    E.g. `value` becomes
+    `_MISSING if (_MAYBE := value) is None else _MAYBE`
+    """
+    _store_MAYBE = ast.Name(id="_MAYBE", ctx=ast.Store())
+    _load_MAYBE = ast.Name(id="_MAYBE", ctx=ast.Load())
+    return ast.IfExp(
+        test=ast.Compare(
+            left=ast.NamedExpr(target=_store_MAYBE, value=value),
+            ops=[ast.Is()],
+            comparators=[ast.Constant(value=None)],
+        ),
+        body=ast.Name(id="_MISSING", ctx=ast.Load()),
+        orelse=_as_str(_load_MAYBE) if as_str else _load_MAYBE,
+    )
+
+
+def _as_str(value: ast.expr) -> ast.Call:
+    """E.g. `value` becomes `str(value)`"""
+    return ast.Call(func=ast.Name(id="str", ctx=ast.Load()), args=[value], keywords=[])
+
+
+def _ast_replacement(expression: str) -> ast.expr:
+    as_str = expression in (
+        "span.status_code",
+        "span.span_kind",
+        "span.parent_id",
+        "span.context.span_id",
+        "span.context.trace_id",
+    )
+    return _replace_none_with_missing(ast.parse(expression, mode="eval").body, as_str)
+
+
+def _allowed_replacements() -> Iterator[Tuple[str, ast.expr]]:
+    for source_segment, ast_replacement in {
+        "name": _ast_replacement("span.name"),
+        "status_code": _ast_replacement("span.status_code"),
+        "span_kind": _ast_replacement("span.span_kind"),
+        "parent_id": _ast_replacement("span.parent_id"),
+    }.items():
+        yield source_segment, ast_replacement
+        yield "span." + source_segment, ast_replacement
+    for source_segment, ast_replacement in {
+        "span_id": _ast_replacement("span.context.span_id"),
+        "trace_id": _ast_replacement("span.context.trace_id"),
+    }.items():
+        yield source_segment, ast_replacement
+        yield "context." + source_segment, ast_replacement
+        yield "span.context." + source_segment, ast_replacement
+    for source_segment, ast_replacement in {
+        field_name: _ast_replacement(f"span.attributes.get('{field_name}')")
+        for field_name in (
+            getattr(semantic_conventions, variable_name)
+            for variable_name in dir(semantic_conventions)
+            if variable_name.isupper()
+            and variable_name.startswith(("RETRIEVAL", "EMBEDDING", "LLM", "TOOL"))
+        )
+    }.items():
+        yield source_segment, ast_replacement
+        yield "attributes." + source_segment, ast_replacement
+        yield "span.attributes." + source_segment, ast_replacement
+
+
+class _Translator(ast.NodeTransformer):
+    _allowed_fields: Mapping[str, ast.expr] = dict(_allowed_replacements())
+
+    def __init__(self, source: str) -> None:
+        # In Python 3.8, we have to use `ast.get_source_segment(source, node)`.
+        # In Python 3.9, we can use `ast.unparse(node)` instead.
+        self._source = source
+
+    def visit_Attribute(self, node: ast.Attribute) -> Any:
+        source_segment: str = cast(str, ast.get_source_segment(self._source, node))
+        if replacement := self._allowed_fields.get(source_segment):
+            return replacement
+        raise SyntaxError(f"invalid expression: {source_segment}")  # TODO: add details
+
+    def visit_Name(self, node: ast.Name) -> Any:
+        source_segment: str = cast(str, ast.get_source_segment(self._source, node))
+        if replacement := self._allowed_fields.get(source_segment):
+            return replacement
+        raise SyntaxError(f"invalid expression: {source_segment}")  # TODO: add details
+
+    def visit_Constant(self, node: ast.Constant) -> Any:
+        return ast.Name(id="_MISSING", ctx=ast.Load()) if node.value is None else node
+
+
+def _validate_expression(expression: ast.Expression, source: str) -> None:
+    # In Python 3.8, we have to use `ast.get_source_segment(source, node)`.
+    # In Python 3.9, we can use `ast.unparse(node)` instead.
+    if not isinstance(expression, ast.Expression):
+        raise SyntaxError(f"invalid expression: {source}")  # TODO: add details
+    for i, node in enumerate(ast.walk(expression.body)):
+        if i == 0:
+            if isinstance(node, (ast.BoolOp, ast.Compare)):
+                continue
+        elif isinstance(
+            node,
+            (
+                ast.Attribute,
+                ast.BinOp,
+                ast.BoolOp,
+                ast.Compare,
+                ast.Constant,
+                ast.Load,
+                ast.Name,
+                ast.Tuple,
+                ast.List,
+                ast.UnaryOp,
+                ast.boolop,
+                ast.cmpop,
+                ast.operator,
+                ast.unaryop,
+            ),
+        ):
+            continue
+        source_segment = cast(str, ast.get_source_segment(source, node))
+        raise SyntaxError(f"invalid expression: {source_segment}")  # TODO: add details
+
+
+class _Missing:
+    """Falsifies all comparisons except those with self."""
+
+    def __lt__(self, other: Any) -> bool:
+        return False
+
+    def __le__(self, other: Any) -> bool:
+        return False
+
+    def __gt__(self, other: Any) -> bool:
+        return False
+
+    def __ge__(self, other: Any) -> bool:
+        return False
+
+    def __eq__(self, other: Any) -> bool:
+        return isinstance(other, _Missing)
+
+    def __ne__(self, other: Any) -> bool:
+        return False
+
+    def __len__(self) -> int:
+        return 0
+
+    def __iter__(self) -> Any:
+        return self
+
+    def __next__(self) -> Any:
+        raise StopIteration()
+
+    def __contains__(self, item: Any) -> bool:
+        return False
+
+    def __str__(self) -> str:
+        return ""
+
+    def __float__(self) -> float:
+        return float("nan")