feat: Add function to upload to trace viewer to Tracer. (#870)

* 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 <[email protected]> Co-authored-by: FelixFehse <[email protected]> Co-authored-by: Sebastian Niehus <[email protected]>
Aleph-Alpha · Jun 4, 2024 · 725c713 · 725c713
1 parent 9ff1d21
commit 725c713
Show file tree

Hide file tree

Showing 7 changed files with 89 additions and 41 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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")