From 725c713c21ed587780459e0b230fc0c36ac60b02 Mon Sep 17 00:00:00 2001 From: FlorianSchepersAA <163116895+FlorianSchepersAA@users.noreply.github.com> Date: Tue, 4 Jun 2024 14:34:52 +0200 Subject: [PATCH] feat: Add function to upload to trace viewer to Tracer. (#870) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: Add function to upload to trace viewer to Tracer. * fix: improve test for trace viewer submission * fix: IL-498 correct json dumping of trace * feat: add file tracer file conversion to new format TASK: IL-496 --------- Co-authored-by: Niklas Köhnecke Co-authored-by: FelixFehse Co-authored-by: Sebastian Niehus --- .github/workflows/sdk-tests.yml | 7 ++++ CHANGELOG.md | 2 + docker-compose.yaml | 8 +++- .../core/tracer/file_tracer.py | 8 ++++ .../core/tracer/in_memory_tracer.py | 29 +------------- src/intelligence_layer/core/tracer/tracer.py | 37 ++++++++++++++++-- tests/core/tracer/test_tracer.py | 39 ++++++++++++++----- 7 files changed, 89 insertions(+), 41 deletions(-) diff --git a/.github/workflows/sdk-tests.yml b/.github/workflows/sdk-tests.yml index 7a0629b09..a4ac3263e 100644 --- a/.github/workflows/sdk-tests.yml +++ b/.github/workflows/sdk-tests.yml @@ -122,6 +122,13 @@ jobs: - "4317:4317" - "4318:4318" - "16686:16686" +# trace-viewer: +# image: ghcr.io/aleph-alpha/trace-viewer-trace-viewer:main +# credentials: +# username: "unused" +# password: ${{ secrets.GITHUB_TOKEN }} # TODO: add PAT +# ports: +# - "3000:3000" steps: - name: Checkout repository uses: actions/checkout@v4 diff --git a/CHANGELOG.md b/CHANGELOG.md index f3e481ebc..5261067f4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,8 @@ - All models raise an error during initialization if an incompatible `name` is passed, instead of only when they are used. - Add `aggregation_overviews_to_pandas` function to allow for easier comparison of multiple aggregation overviews - Add `parameter_optimization.ipynb` notebook to demonstrate the optimization of tasks by comparing different parameter combinations. + - Add `convert_file_for_viewing` in the `FileTracer` to convert the trace file format to the new (OpenTelemetry style) format and save as a new file. + - All tracers can now call `submit_to_trace_viewer` to send the trace to the Trace Viewer. ### Fixes - The document index client now correctly URL-encodes document names in its queries. diff --git a/docker-compose.yaml b/docker-compose.yaml index 8d288a0c9..b01e865fa 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -13,7 +13,6 @@ services: environment: ARGILLA_ELASTICSEARCH: "http://argilla-elastic-search:9200" ARGILLA_ENABLE_TELEMETRY: 0 - open-telemetry-trace-service: container_name: jaeger_1_35 environment: @@ -23,3 +22,10 @@ services: - "4318:4318" - "16686:16686" image: jaegertracing/all-in-one:1.35 + # TODO: Below code should be enabled once the secret has been added to the CI. After that, also check tests in test_tracer.py + # export GITHUB_TOKEN=... + # echo $GITHUB_TOKEN | docker login ghcr.io -u your_email@for_github --password-stdin + # trace_viewer: + # image: ghcr.io/aleph-alpha/trace-viewer-trace-viewer:main + # ports: + # - 3000:3000 diff --git a/src/intelligence_layer/core/tracer/file_tracer.py b/src/intelligence_layer/core/tracer/file_tracer.py index 5c7813caa..10dbb6d77 100644 --- a/src/intelligence_layer/core/tracer/file_tracer.py +++ b/src/intelligence_layer/core/tracer/file_tracer.py @@ -78,6 +78,14 @@ def traces(self, trace_id: Optional[str] = None) -> InMemoryTracer: ) return self._parse_log(filtered_traces) + def convert_file_for_viewing(self, file_path: Path | str) -> None: + in_memory_tracer = self.traces() + traces = in_memory_tracer.export_for_viewing() + path_to_file = Path(file_path) + with path_to_file.open(mode="w", encoding="utf-8") as file: + for exportedSpan in traces: + file.write(exportedSpan.model_dump_json() + "\n") + class FileSpan(PersistentSpan, FileTracer): """A `Span` created by `FileTracer.span`.""" diff --git a/src/intelligence_layer/core/tracer/in_memory_tracer.py b/src/intelligence_layer/core/tracer/in_memory_tracer.py index 534b1b440..62c932d1a 100644 --- a/src/intelligence_layer/core/tracer/in_memory_tracer.py +++ b/src/intelligence_layer/core/tracer/in_memory_tracer.py @@ -1,12 +1,9 @@ -import os from datetime import datetime from typing import Optional, Sequence, Union from uuid import UUID -import requests import rich from pydantic import BaseModel, Field, SerializeAsAny -from requests import HTTPError from rich.panel import Panel from rich.syntax import Syntax from rich.tree import Tree @@ -15,7 +12,6 @@ Context, Event, ExportedSpan, - ExportedSpanList, JsonSerializer, PydanticSerializable, Span, @@ -81,27 +77,6 @@ def _ipython_display_(self) -> None: if not self.submit_to_trace_viewer(): rich.print(self._rich_render_()) - def submit_to_trace_viewer(self) -> bool: - """Submits the trace to the UI for visualization""" - trace_viewer_url = os.getenv("TRACE_VIEWER_URL", "http://localhost:3000") - trace_viewer_trace_upload = f"{trace_viewer_url}/trace" - try: - res = requests.post( - trace_viewer_trace_upload, - json=ExportedSpanList(self.export_for_viewing()).model_dump_json(), - ) - if res.status_code != 200: - raise HTTPError(res.status_code) - rich.print( - f"Open the [link={trace_viewer_url}]Trace Viewer[/link] to view the trace." - ) - return True - except requests.ConnectionError: - print( - f"Trace viewer not found under {trace_viewer_url}.\nConsider running it for a better viewing experience.\nIf it is, set `TRACE_VIEWER_URL` in the environment." - ) - return False - def export_for_viewing(self) -> Sequence[ExportedSpan]: exported_root_spans: list[ExportedSpan] = [] for entry in self.entries: @@ -177,7 +152,7 @@ def _rich_render_(self) -> Tree: return tree - def _span_attributes(self) -> SpanAttributes: + def _span_attributes(self) -> SpanAttributes | TaskSpanAttributes: return SpanAttributes() def export_for_viewing(self) -> Sequence[ExportedSpan]: @@ -253,7 +228,7 @@ def __init__( def record_output(self, output: PydanticSerializable) -> None: self.output = output - def _span_attributes(self) -> SpanAttributes: + def _span_attributes(self) -> SpanAttributes | TaskSpanAttributes: return TaskSpanAttributes(input=self.input, output=self.output) def _rich_render_(self) -> Tree: diff --git a/src/intelligence_layer/core/tracer/tracer.py b/src/intelligence_layer/core/tracer/tracer.py index f97a6c9e4..e6ea1b02e 100644 --- a/src/intelligence_layer/core/tracer/tracer.py +++ b/src/intelligence_layer/core/tracer/tracer.py @@ -1,12 +1,15 @@ +import os import traceback from abc import ABC, abstractmethod from contextlib import AbstractContextManager from datetime import datetime, timezone from enum import Enum from types import TracebackType -from typing import TYPE_CHECKING, Mapping, Optional, Sequence +from typing import TYPE_CHECKING, Mapping, Optional, Sequence, Union from uuid import UUID, uuid4 +import requests +import rich from pydantic import BaseModel, Field, RootModel, SerializeAsAny from typing_extensions import Self, TypeAliasType @@ -65,7 +68,7 @@ class SpanAttributes(BaseModel): type: SpanType = SpanType.SPAN -class TaskSpanAttributes(SpanAttributes): +class TaskSpanAttributes(BaseModel): type: SpanType = SpanType.TASK_SPAN input: SerializeAsAny[PydanticSerializable] output: SerializeAsAny[PydanticSerializable] @@ -87,7 +90,7 @@ class ExportedSpan(BaseModel): parent_id: UUID | None start_time: datetime end_time: datetime - attributes: SpanAttributes + attributes: Union[SpanAttributes, TaskSpanAttributes] events: Sequence[Event] status: SpanStatus # we ignore the links concept @@ -170,6 +173,34 @@ def export_for_viewing(self) -> Sequence[ExportedSpan]: """ ... + def submit_to_trace_viewer(self) -> bool: + """Submits the trace to the UI for visualization""" + trace_viewer_url = os.getenv("TRACE_VIEWER_URL", "http://localhost:3000") + trace_viewer_trace_upload = f"{trace_viewer_url}/trace" + try: + res = requests.post( + trace_viewer_trace_upload, + headers={ + "Content-Type": "application/json", + "Accept": "application/json", + }, + json=ExportedSpanList(self.export_for_viewing()).model_dump( + mode="json" + ), + ) + print(res) + if res.status_code != 200: + raise requests.HTTPError(res.status_code) + rich.print( + f"Open the [link={trace_viewer_url}]Trace Viewer[/link] to view the trace." + ) + return True + except requests.ConnectionError: + print( + f"Trace viewer not found under {trace_viewer_url}.\nConsider running it for a better viewing experience.\nIf it is, set `TRACE_VIEWER_URL` in the environment." + ) + return False + class ErrorValue(BaseModel): error_type: str diff --git a/tests/core/tracer/test_tracer.py b/tests/core/tracer/test_tracer.py index dfffc59f6..3cd154782 100644 --- a/tests/core/tracer/test_tracer.py +++ b/tests/core/tracer/test_tracer.py @@ -14,6 +14,7 @@ Tracer, utc_now, ) +from intelligence_layer.core.task import Task from tests.core.tracer.conftest import SpecificTestException @@ -207,38 +208,56 @@ def test_tracer_raises_if_open_span_is_exported( child_span.export_for_viewing() -@pytest.mark.skip("Not yet implemented") @pytest.mark.parametrize( "tracer_fixture", tracer_fixtures, ) -def test_spans_cannot_be_closed_twice( +def test_spans_cannot_be_used_as_context_twice( tracer_fixture: str, request: pytest.FixtureRequest, ) -> None: tracer: Tracer = request.getfixturevalue(tracer_fixture) span = tracer.span("name") - span.end() - span.end() + with span: + pass + with pytest.raises(Exception): + with span: + pass +@pytest.mark.skip("Not yet implemented") +@pytest.mark.docker @pytest.mark.parametrize( "tracer_fixture", tracer_fixtures, ) -def test_spans_cannot_be_used_as_context_twice( +def test_tracer_can_be_submitted_to_trace_viewer( + tracer_fixture: str, + request: pytest.FixtureRequest, + tracer_test_task: Task[str, str], +) -> None: + tracer: Tracer = request.getfixturevalue(tracer_fixture) + + tracer_test_task.run(input="input", tracer=tracer) + + assert tracer.submit_to_trace_viewer() + + +@pytest.mark.skip("Not yet implemented") +@pytest.mark.parametrize( + "tracer_fixture", + tracer_fixtures, +) +def test_spans_cannot_be_closed_twice( tracer_fixture: str, request: pytest.FixtureRequest, ) -> None: tracer: Tracer = request.getfixturevalue(tracer_fixture) span = tracer.span("name") - with span: - pass - with pytest.raises(Exception): - with span: - pass + span.end() + span.end() @pytest.mark.skip("Not yet implemented")