Improve profiler scheduling

Author: dmelikyan

graphsignal/__init__.py

@@ -25,7 +25,7 @@ def configure(
     deployment: Optional[str] = None,
     tags: Optional[Dict[str, str]] = None,
     auto_instrument: Optional[bool] = None,
-    profiling_rate: Optional[float] = None,
+    profiles_per_min: Optional[int] = None,
     include_profiles: Optional[list] = None,
     debug_mode: Optional[bool] = None
 ) -> None:
@@ -39,7 +39,7 @@ def configure(
     api_url = read_config_param("api_url", str, api_url)
     tags = read_config_tags(tags)
     auto_instrument = read_config_param("auto_instrument", bool, auto_instrument, default_value=True)
-    profiling_rate = read_config_param("profiling_rate", float, profiling_rate, default_value=0.1)
+    profiles_per_min = read_config_param("profiles_per_min", int, profiles_per_min)
     include_profiles = read_config_param("include_profiles", list, include_profiles)
     debug_mode = read_config_param("debug_mode", bool, debug_mode, default_value=False)
 
@@ -52,7 +52,7 @@ def configure(
         api_url=api_url,
         tags=tags,
         auto_instrument=auto_instrument,
-        profiling_rate=profiling_rate,
+        profiles_per_min=profiles_per_min,
         include_profiles=include_profiles,
         debug_mode=debug_mode)
     _tracer.setup()

graphsignal/profiles.py

@@ -20,9 +20,8 @@ def _can_include_profiles(self, span, profiles):
                 span.can_include_profiles(profiles))
 
     def on_span_start(self, span, context):
-        if (span.sampled() and 
-            self._can_include_profiles(span, ['profile.cpython']) and 
-            graphsignal._tracer.set_profiling_mode()):
+        if (self._can_include_profiles(span, ['profile.cpython']) and 
+            graphsignal._tracer.set_profiling_mode('profile.cpython')):
 
             context['profiled'] = True
 

graphsignal/recorders/python_recorder.py

@@ -25,9 +25,8 @@ def _can_include_profiles(self, span, profiles):
                 span.can_include_profiles(profiles))
 
     def on_span_start(self, span, context):
-        if (span.sampled() and 
-            self._can_include_profiles(span, ['profile.pytorch']) and 
-            graphsignal._tracer.set_profiling_mode()):
+        if (self._can_include_profiles(span, ['profile.pytorch']) and 
+            graphsignal._tracer.set_profiling_mode('profile.pytorch')):
             context['profiled'] = True
 
             if self._torch_prof:

graphsignal/recorders/pytorch_recorder.py

@@ -52,14 +52,12 @@ def __init__(self, name, value, unit=None):
 class SpanContext:
     __slots__ = [
         'trace_id',
-        'span_id',
-        'sampled'
+        'span_id'
     ]
 
-    def __init__(self, trace_id=None, span_id=None, sampled=None):
+    def __init__(self, trace_id=None, span_id=None):
         self.trace_id = trace_id
         self.span_id = span_id
-        self.sampled = sampled
 
     @staticmethod
     def push_contextvars(ctx):
@@ -80,26 +78,24 @@ def loads(value):
                 logger.debug(f'SpanContext.loads: invalid context value: {value}')
             return None
         parts = value.split('-')
-        if len(parts) < 3:
+        if len(parts) < 2:
             if logger.isEnabledFor(logging.DEBUG):
                 logger.debug(f'SpanContext.loads: invalid context value: {value}')
             return None
         ctx = SpanContext()
         ctx.trace_id = parts[0]
         ctx.span_id = parts[1]
-        ctx.sampled = parts[2] == '1'
         return ctx
 
     @staticmethod
     def dumps(ctx):
-        if ctx is None or ctx.trace_id is None or ctx.span_id is None or ctx.sampled is None:
+        if ctx is None or ctx.trace_id is None or ctx.span_id is None:
             if logger.isEnabledFor(logging.DEBUG):
                 logger.debug('SpanContext.dumps: invalid context')
             return None
-        return '{0}-{1}-{2}'.format(
+        return '{0}-{1}'.format(
             ctx.trace_id, 
-            ctx.span_id, 
-            '1' if ctx.sampled else '0')
+            ctx.span_id)
 
 class Span:
     MAX_SPAN_TAGS = 25
@@ -120,7 +116,6 @@ class Span:
         '_parent_span_id',
         '_linked_span_ids',
         '_is_root',
-        '_sampled',
         '_recorder_context',
         '_model',
         '_is_started',
@@ -161,11 +156,9 @@ def __init__(self, operation, tags=None, include_profiles=None, parent_context=N
         self._span_id = None
         self._trace_id = None
         self._parent_span_id = None
-        self._sampled = None
         if parent_context:
             self._trace_id = parent_context.trace_id
             self._parent_span_id = parent_context.span_id
-            self._sampled = parent_context.sampled
         self._linked_span_ids = None
         self._is_root = False
         self._start_us = None
@@ -209,6 +202,9 @@ async def __aexit__(self, *exc_info):
         self.stop()
         return False
 
+    def _should_record(self) -> bool:
+        return self._profiles is not None and len(self._profiles) > 0
+
     def _start(self):
         if self._is_started:
             return
@@ -222,9 +218,7 @@ def _start(self):
         if self._trace_id is None:
             self._trace_id = uuid_sha1(size=12)
             self._is_root = True
-        if self._sampled is None:
-            self._sampled = fast_rand() <= graphsignal._tracer.sampling_rate
-
+        
         self._context_tags = _tracer().context_tags.get().copy()
 
         self._model = client.Span(
@@ -307,7 +301,16 @@ def _stop(self) -> None:
                 _tracer().metric_store().inc_counter(
                     name=metric.name, tags=span_tags, value=metric.value, update_ts=now)
 
-        if self._sampled:
+        # report issues
+        if self._exc_infos:
+            for exc_info in self._exc_infos:
+                _tracer().report_issue(
+                    name='operation.error',
+                    severity=1,
+                    description=f'{exc_info[0].__name__}: {exc_info[1]}',
+                    span=self)
+
+        if self._should_record():
             # fill and upload span
             # copy data to span model
             self._model.start_us = self._start_us
@@ -404,8 +407,7 @@ def trace_id(self) -> str:
     def get_span_context(self):
         return SpanContext(
             trace_id=self._trace_id,
-            span_id=self._span_id,
-            sampled=self._sampled)
+            span_id=self._span_id)
 
     def add_linked_span(self, span_id: str) -> None:
         if not span_id:
@@ -417,9 +419,6 @@ def add_linked_span(self, span_id: str) -> None:
 
         self._linked_span_ids.append(span_id)
 
-    def sampled(self) -> bool:
-        return self._sampled
-
     def can_include_profiles(self, profiles) -> bool:
         if self._include_profiles_index is None:
             return True

graphsignal/spans.py

@@ -22,6 +22,7 @@
 logger = logging.getLogger('graphsignal')
 
 
+
 class GraphsignalTracerLogHandler(logging.Handler):
     def __init__(self, tracer):
         super().__init__()
@@ -102,6 +103,34 @@ def find_spec(self, fullname, path=None, target=None):
 
         return None
 
+class ProfilingTokenBucket:
+    def __init__(self, sampling_rate_per_minute: float):
+        self.capacity = sampling_rate_per_minute   # max tokens (samples) per minute
+        self.tokens = self.capacity               # start full
+        self.refill_rate_per_sec = self.capacity / 60.0  # tokens per second
+        self.last_refill_time = time.monotonic()
+        self._first_request_skipped = False  # Track if first request has been skipped
+    
+    def _refill(self):
+        now = time.monotonic()
+        elapsed = now - self.last_refill_time
+        if elapsed > 0:
+            # Add tokens based on elapsed time
+            new_tokens = elapsed * self.refill_rate_per_sec
+            self.tokens = min(self.capacity, self.tokens + new_tokens)
+            self.last_refill_time = now
+    
+    def should_profile(self) -> bool:
+        # Skip the first profiling request
+        if not self._first_request_skipped:
+            self._first_request_skipped = True
+            return False
+        
+        self._refill()
+        if self.tokens >= 1:
+            self.tokens -= 1
+            return True
+        return False
 
 class Tracer:
     TICK_INTERVAL_SEC = 10
@@ -114,8 +143,7 @@ def __init__(
             api_url=None, 
             tags=None, 
             auto_instrument=True, 
-            sampling_rate=1.0,
-            profiling_rate=0.1,
+            profiles_per_min=10,
             include_profiles=None,
             debug_mode=False):
         if debug_mode:
@@ -139,11 +167,12 @@ def __init__(
         self.params = {}
 
         self.auto_instrument = auto_instrument
-        self.sampling_rate = sampling_rate if sampling_rate is not None else 1.0
-        self.profiling_rate = profiling_rate if profiling_rate is not None else 0
+        self.profiles_per_min = profiles_per_min if profiles_per_min is not None else 10
         self.include_profiles = include_profiles
         self.debug_mode = debug_mode
 
+        self._profiling_token_buckets = {}
+
         self._tick_timer_thread = None
         self._tick_stop_event = threading.Event()
         self._tick_lock = threading.Lock()
@@ -272,8 +301,11 @@ def metric_store(self):
     def log_store(self):
         return self._log_store
 
-    def set_profiling_mode(self):
-        if fast_rand() > self.profiling_rate:
+    def set_profiling_mode(self, profile_name):
+        if profile_name not in self._profiling_token_buckets:
+            self._profiling_token_buckets[profile_name] = ProfilingTokenBucket(self.profiles_per_min)
+
+        if not self._profiling_token_buckets[profile_name].should_profile():
             return False
 
         with self._profiling_mode_lock:

graphsignal/tracer.py

@@ -1 +1 @@
-__version__ = '0.18.0'
+__version__ = '0.18.1'

graphsignal/version.py

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphsignal"
-version = "0.18.0"
+version = "0.18.1"
 description = "Graphsignal Tracer for Python"
 authors = ["Graphsignal, Inc. <devops@graphsignal.com>"]
 license = "Apache-2.0"

pyproject.toml

@@ -11,6 +11,7 @@
 
 import graphsignal
 from graphsignal.recorders.instrumentation import patch_method, trace_method, profile_method, read_args, parse_semver, compare_semver
+from graphsignal.spans import Span
 from graphsignal.uploader import Uploader
 from test.model_utils import find_tag, find_counter
 
@@ -98,8 +99,9 @@ def profile_func(event_name, duration_ns):
         self.assertTrue(measured_duration_ns > 0)
 
 
+    @patch.object(Span, '_should_record', return_value=True)
     @patch.object(Uploader, 'upload_span')
-    async def test_trace_method(self, mocked_upload_span):
+    async def test_trace_method(self, mocked_upload_span, mocked_should_record):
         obj = Dummy()
 
         trace_func_called = False
@@ -116,8 +118,9 @@ def trace_func(span, args, kwargs, ret, exc):
         self.assertTrue(trace_func_called)
         self.assertEqual(find_tag(model, 'operation.name'), 'op1')
 
+    @patch.object(Span, '_should_record', return_value=True)
     @patch.object(Uploader, 'upload_span')
-    async def test_trace_method_generator(self, mocked_upload_span):
+    async def test_trace_method_generator(self, mocked_upload_span, mocked_should_record):
         obj = Dummy()
 
         trace_func_called = None

test/recorders/test_instrumentation.py

@@ -23,7 +23,6 @@ def setUp(self):
             logger.addHandler(logging.StreamHandler(sys.stdout))
         graphsignal.configure(
             api_key='k1',
-            profiling_rate=1,
             debug_mode=True)
         graphsignal._tracer.auto_export = False
 
@@ -34,6 +33,10 @@ def tearDown(self):
     @patch.object(Tracer, 'emit_span_stop')
     @patch.object(Tracer, 'emit_span_read')
     def test_record(self, mocked_emit_span_read, mocked_emit_span_stop, mocked_emit_span_start):
+        # First call will be skipped, second call should succeed
+        graphsignal._tracer.set_profiling_mode('profile.cpython')
+        graphsignal._tracer.unset_profiling_mode()
+
         recorder = PythonRecorder()
         recorder._exclude_path = 'donotmatchpath'
         recorder.setup()