Rework trace_id

Author: dmelikyan

graphsignal/__init__.py

@@ -5,7 +5,7 @@
 
 from graphsignal.version import __version__
 from graphsignal.tracer import Tracer
-from graphsignal.spans import Span
+from graphsignal.spans import Span, SpanContext
 
 logger = logging.getLogger('graphsignal')
 
@@ -210,6 +210,7 @@ def shutdown() -> None:
     'shutdown',
     'trace',
     'function_trace',
+    'SpanContext',
     'Span',
     'score',
     'set_tag',

graphsignal/client/models/span.py

@@ -32,7 +32,7 @@ class Span(BaseModel):
     Span
     """ # noqa: E501
     span_id: StrictStr = Field(description="Unique identifier for the span.")
-    root_span_id: Optional[StrictStr] = Field(default=None, description="Identifier of the root span, if this is a child span.")
+    trace_id: Optional[StrictStr] = Field(default=None, description="Identifier of the trace.")
     parent_span_id: Optional[StrictStr] = Field(default=None, description="Identifier of the parent span, if this is a nested span.")
     linked_span_ids: Optional[List[StrictStr]] = Field(default=None, description="List of linked span identifiers.")
     start_us: StrictInt = Field(description="Start time of the span in microseconds.")
@@ -42,7 +42,7 @@ class Span(BaseModel):
     params: Optional[List[Param]] = Field(default=None, description="List of parameters associated with the span.")
     counters: Optional[List[Counter]] = Field(default=None, description="Counters associated with the span.")
     profiles: Optional[List[Profile]] = Field(default=None, description="List of profiles related to the span.")
-    __properties: ClassVar[List[str]] = ["span_id", "root_span_id", "parent_span_id", "linked_span_ids", "start_us", "end_us", "tags", "exceptions", "params", "counters", "profiles"]
+    __properties: ClassVar[List[str]] = ["span_id", "trace_id", "parent_span_id", "linked_span_ids", "start_us", "end_us", "tags", "exceptions", "params", "counters", "profiles"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -131,7 +131,7 @@ def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
 
         _obj = cls.model_validate({
             "span_id": obj.get("span_id"),
-            "root_span_id": obj.get("root_span_id"),
+            "trace_id": obj.get("trace_id"),
             "parent_span_id": obj.get("parent_span_id"),
             "linked_span_ids": obj.get("linked_span_ids"),
             "start_us": obj.get("start_us"),

graphsignal/recorders/pytorch_recorder.py

@@ -1,6 +1,5 @@
 import logging
 import os
-import random
 import json
 import time
 import torch
@@ -26,7 +25,7 @@ def on_span_start(self, span, context):
             # profiler active, skip
             # only one profiler is allowed per process, so we don't need global locks
             return
-        if random.random() > graphsignal._tracer.profiling_rate:
+        if not span.profiled():
             return
 
         if not self._torch_prof: