Fix SENTRY_TRACES_BACKGROUND_SAMPLE_RATE not respected

r4victor · r4victor · commit 06b444aa7406 · 2026-04-06T09:49:51.000+05:00
diff --git a/src/dstack/_internal/server/app.py b/src/dstack/_internal/server/app.py
@@ -14,7 +14,6 @@
 from fastapi.staticfiles import StaticFiles
 from packaging.version import Version
 from prometheus_client import Counter, Histogram
-from sentry_sdk.types import SamplingContext
 
 from dstack._internal import settings as core_settings
 from dstack._internal.cli.utils.common import console
@@ -115,7 +114,7 @@ async def lifespan(app: FastAPI):
             release=core_settings.DSTACK_VERSION,
             environment=settings.SERVER_ENVIRONMENT,
             enable_tracing=True,
-            traces_sampler=_sentry_traces_sampler,
+            traces_sampler=sentry_utils.sentry_traces_sampler,
             profiles_sample_rate=settings.SENTRY_PROFILES_SAMPLE_RATE,
             before_send=sentry_utils.AsyncioCancelledErrorFilterEventProcessor(),
         )
@@ -426,18 +425,6 @@ def _is_prometheus_request(request: Request) -> bool:
     return request.url.path.startswith("/metrics")
 
 
-def _sentry_traces_sampler(sampling_context: SamplingContext) -> float:
-    parent_sampling_decision = sampling_context["parent_sampled"]
-    if parent_sampling_decision is not None:
-        return float(parent_sampling_decision)
-    transaction_context = sampling_context["transaction_context"]
-    name = transaction_context.get("name")
-    if name is not None:
-        if name.startswith("background."):
-            return settings.SENTRY_TRACES_BACKGROUND_SAMPLE_RATE
-    return settings.SENTRY_TRACES_SAMPLE_RATE
-
-
 def _print_dstack_logo():
     console.print(
         r"""[purple]     _     _             _
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/compute_groups.py b/src/dstack/_internal/server/background/pipeline_tasks/compute_groups.py
@@ -120,7 +120,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.ComputeGroupFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("ComputeGroupFetcher.fetch")
     async def fetch(self, limit: int) -> list[PipelineItem]:
         compute_group_lock, _ = get_locker(get_db().dialect_name).get_lockset(
             ComputeGroupModel.__tablename__
@@ -188,7 +188,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.ComputeGroupWorker.process")
+    @sentry_utils.instrument_pipeline_task("ComputeGroupWorker.process")
     async def process(self, item: PipelineItem):
         async with get_session_ctx() as session:
             res = await session.execute(
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/fleets.py b/src/dstack/_internal/server/background/pipeline_tasks/fleets.py
@@ -134,7 +134,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.FleetFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("FleetFetcher.fetch")
     async def fetch(self, limit: int) -> list[PipelineItem]:
         fleet_lock, _ = get_locker(get_db().dialect_name).get_lockset(FleetModel.__tablename__)
         async with fleet_lock:
@@ -203,7 +203,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.FleetWorker.process")
+    @sentry_utils.instrument_pipeline_task("FleetWorker.process")
     async def process(self, item: PipelineItem):
         process_context = await _load_process_context(item)
         if process_context is None:
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py
@@ -129,7 +129,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.GatewayFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("GatewayFetcher.fetch")
     async def fetch(self, limit: int) -> list[GatewayPipelineItem]:
         gateway_lock, _ = get_locker(get_db().dialect_name).get_lockset(GatewayModel.__tablename__)
         async with gateway_lock:
@@ -207,7 +207,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.GatewayWorker.process")
+    @sentry_utils.instrument_pipeline_task("GatewayWorker.process")
     async def process(self, item: GatewayPipelineItem):
         if item.to_be_deleted:
             await _process_to_be_deleted_item(item)
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/__init__.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/__init__.py
@@ -152,7 +152,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.InstanceFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("InstanceFetcher.fetch")
     async def fetch(self, limit: int) -> list[InstancePipelineItem]:
         instance_lock, _ = get_locker(get_db().dialect_name).get_lockset(
             InstanceModel.__tablename__
@@ -267,7 +267,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.InstanceWorker.process")
+    @sentry_utils.instrument_pipeline_task("InstanceWorker.process")
     async def process(self, item: InstancePipelineItem):
         process_context: Optional[_ProcessContext] = None
         if item.status == InstanceStatus.PENDING:
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py
@@ -194,7 +194,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.JobRunningFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("JobRunningFetcher.fetch")
     async def fetch(self, limit: int) -> list[JobRunningPipelineItem]:
         job_lock, _ = get_locker(get_db().dialect_name).get_lockset(JobModel.__tablename__)
         async with job_lock:
@@ -286,7 +286,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.JobRunningWorker.process")
+    @sentry_utils.instrument_pipeline_task("JobRunningWorker.process")
     async def process(self, item: JobRunningPipelineItem):
         context = await _load_process_context(item=item)
         if context is None:
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py
@@ -227,7 +227,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.JobSubmittedFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("JobSubmittedFetcher.fetch")
     async def fetch(self, limit: int) -> list[JobSubmittedPipelineItem]:
         now = get_current_datetime()
         if limit <= 0:
@@ -309,7 +309,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.JobSubmittedWorker.process")
+    @sentry_utils.instrument_pipeline_task("JobSubmittedWorker.process")
     async def process(self, item: JobSubmittedPipelineItem):
         context = await _load_process_context(item=item)
         if context is None:
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py
@@ -160,7 +160,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.JobTerminatingFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("JobTerminatingFetcher.fetch")
     async def fetch(self, limit: int) -> list[JobTerminatingPipelineItem]:
         job_lock, _ = get_locker(get_db().dialect_name).get_lockset(JobModel.__tablename__)
         async with job_lock:
@@ -243,7 +243,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.JobTerminatingWorker.process")
+    @sentry_utils.instrument_pipeline_task("JobTerminatingWorker.process")
     async def process(self, item: JobTerminatingPipelineItem):
         async with get_session_ctx() as session:
             job_model = await _refetch_locked_job(session=session, item=item)
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/placement_groups.py b/src/dstack/_internal/server/background/pipeline_tasks/placement_groups.py
@@ -117,7 +117,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.PlacementGroupFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("PlacementGroupFetcher.fetch")
     async def fetch(self, limit: int) -> list[PipelineItem]:
         placement_group_lock, _ = get_locker(get_db().dialect_name).get_lockset(
             PlacementGroupModel.__tablename__
@@ -187,7 +187,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.PlacementGroupWorker.process")
+    @sentry_utils.instrument_pipeline_task("PlacementGroupWorker.process")
     async def process(self, item: PipelineItem):
         async with get_session_ctx() as session:
             res = await session.execute(
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py
@@ -129,7 +129,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.RunFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("RunFetcher.fetch")
     async def fetch(self, limit: int) -> list[RunPipelineItem]:
         if limit <= 0:
             return []
@@ -243,7 +243,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.RunWorker.process")
+    @sentry_utils.instrument_pipeline_task("RunWorker.process")
     async def process(self, item: RunPipelineItem):
         # Currently `dstack` supports runs with
         # * one multi-node replica (multi-node tasks)
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/volumes.py b/src/dstack/_internal/server/background/pipeline_tasks/volumes.py
@@ -132,7 +132,7 @@ def __init__(
             queue_check_delay=queue_check_delay,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.VolumeFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("VolumeFetcher.fetch")
     async def fetch(self, limit: int) -> list[VolumePipelineItem]:
         volume_lock, _ = get_locker(get_db().dialect_name).get_lockset(VolumeModel.__tablename__)
         async with volume_lock:
@@ -209,7 +209,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.VolumeWorker.process")
+    @sentry_utils.instrument_pipeline_task("VolumeWorker.process")
     async def process(self, item: VolumePipelineItem):
         volume_model = await _refetch_locked_volume(item)
         if volume_model is None:
diff --git a/src/dstack/_internal/server/utils/sentry_utils.py b/src/dstack/_internal/server/utils/sentry_utils.py
@@ -3,17 +3,20 @@
 from typing import Optional
 
 import sentry_sdk
-from sentry_sdk.types import Event, Hint
+from sentry_sdk.types import Event, Hint, SamplingContext
+
+from dstack._internal.server import settings
+
+SCHEDULED_TASKS_PREFIX = "scheduled_tasks"
+PIPELINE_TASKS_PREFIX = "pipeline_tasks"
 
 
 def instrument_scheduled_task(f):
-    @functools.wraps(f)
-    async def wrapper(*args, **kwargs):
-        with sentry_sdk.isolation_scope():
-            with sentry_sdk.start_transaction(name=f"scheduled_tasks.{f.__name__}"):
-                return await f(*args, **kwargs)
+    return instrument_named_task(f"{SCHEDULED_TASKS_PREFIX}.{f.__name__}")
 
-    return wrapper
+
+def instrument_pipeline_task(name: str):
+    return instrument_named_task(f"{PIPELINE_TASKS_PREFIX}.{name}")
 
 
 def instrument_named_task(name: str):
@@ -29,10 +32,26 @@ async def wrapper(*args, **kwargs):
     return decorator
 
 
+def sentry_traces_sampler(sampling_context: SamplingContext) -> float:
+    parent_sampling_decision = sampling_context["parent_sampled"]
+    if parent_sampling_decision is not None:
+        return float(parent_sampling_decision)
+    transaction_context = sampling_context["transaction_context"]
+    name = transaction_context.get("name")
+    if name is not None:
+        if _is_background_transaction(name):
+            return settings.SENTRY_TRACES_BACKGROUND_SAMPLE_RATE
+    return settings.SENTRY_TRACES_SAMPLE_RATE
+
+
 class AsyncioCancelledErrorFilterEventProcessor:
     # See https://docs.sentry.io/platforms/python/configuration/filtering/#filtering-error-events
     def __call__(self, event: Event, hint: Hint) -> Optional[Event]:
         exc_info = hint.get("exc_info")
         if exc_info and isinstance(exc_info[1], asyncio.CancelledError):
             return None
         return event
+
+
+def _is_background_transaction(name: str) -> bool:
+    return name.startswith(SCHEDULED_TASKS_PREFIX) or name.startswith(PIPELINE_TASKS_PREFIX)
