Introduce Openlayer tracer

Author: gustavocidornelas

openlayer/tracing/__init__.py

@@ -0,0 +1,51 @@
+"""Module with the different Step classes that can be used in a trace."""
+
+import time
+from typing import Any, Dict, Optional
+
+
+class Step:
+    def __init__(
+        self,
+        name: str,
+        inputs: Optional[Any] = None,
+        output: Optional[Any] = None,
+        metadata: Dict[str, any] = {},
+    ) -> None:
+        self.name = name
+        self.inputs = inputs
+        self.output = output
+        self.metadata = metadata
+
+        self.step_type = "user_call"
+        self.start_time = time.time()
+        self.end_time = None
+        self.ground_truth = None
+        self.latency = None
+
+        self.steps = []
+
+    def add_nested_step(self, nested_step: "Step") -> None:
+        """Adds a nested step to the current step."""
+        self.steps.append(nested_step)
+
+    def update_data(self, **kwargs: Any) -> None:
+        """Updates the step data."""
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Dictionary representation of the Step."""
+        return {
+            "name": self.name,
+            "type": self.step_type,
+            "inputs": self.inputs,
+            "output": self.output,
+            "groundTruth": self.ground_truth,
+            "metadata": self.metadata,
+            "steps": [nested_step.to_dict() for nested_step in self.steps],
+            "latency": self.latency,
+            "startTime": self.start_time,
+            "endTime": self.end_time,
+        }

openlayer/tracing/steps.py

@@ -0,0 +1,80 @@
+"""Module with the logic to create and manage traces and steps."""
+
+import inspect
+from typing import Any, Dict, Optional, Generator
+from contextlib import contextmanager
+import contextvars
+from functools import wraps
+
+from . import steps
+from . import traces
+import time
+
+_current_step = contextvars.ContextVar("current_step")
+_current_trace = contextvars.ContextVar("current_trace")
+
+
+@contextmanager
+def create_step(
+    name: str,
+    inputs: Optional[Any] = None,
+    output: Optional[Any] = None,
+    metadata: Dict[str, any] = {},
+) -> Generator[steps.Step, None, None]:
+    """Starts a trace and yields a Step object."""
+    new_step = steps.Step(name=name, inputs=inputs, output=output, metadata=metadata)
+
+    parent_step = _current_step.get(None)
+    is_root_step = parent_step is None
+
+    if parent_step is None:
+        print("Starting a new trace...")
+        current_trace = traces.Trace()
+        _current_trace.set(current_trace)  # Set the current trace in context
+        current_trace.add_step(new_step)
+    else:
+        print(f"Adding step {name} to parent step {parent_step.name}")
+        current_trace = _current_trace.get()
+        parent_step.add_nested_step(new_step)
+
+    token = _current_step.set(new_step)
+
+    try:
+        yield new_step
+    finally:
+        _current_step.reset(token)
+        if is_root_step:
+            print("Ending the trace...")
+            print("-" * 80)
+            print(current_trace.to_dict())
+            print("-" * 80)
+        else:
+            # TODO: stream to Openlayer
+            print(f"Ending step {name}")
+
+
+def trace(*step_args, **step_kwargs):
+    def decorator(func):
+        func_signature = inspect.signature(func)
+
+        @wraps(func)
+        def wrapper(*func_args, **func_kwargs):
+            if step_kwargs.get("name") is None:
+                step_kwargs["name"] = func.__name__
+            with create_step(*step_args, **step_kwargs) as step:
+                output = func(*func_args, **func_kwargs)
+                end_time = time.time()
+                latency = (end_time - step.start_time) * 1000  # in ms
+                inputs = func_signature.bind(*func_args, **func_kwargs).arguments
+
+                step.update_data(
+                    inputs=inputs,
+                    output=output,
+                    end_time=end_time,
+                    latency=latency,
+                )
+            return output
+
+        return wrapper
+
+    return decorator

openlayer/tracing/tracer.py

@@ -0,0 +1,19 @@
+"""Module with the Trace class."""
+
+from typing import Any, Dict
+
+from .steps import Step
+
+
+class Trace:
+    def __init__(self):
+        self.steps = []
+        self.current_step = None
+
+    def add_step(self, step: Step) -> None:
+        """Adds a step to the trace."""
+        self.steps.append(step)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Dictionary representation of the Trace."""
+        return {"rows": [step.to_dict() for step in self.steps]}