feat(dynamic-sampling): Add per-project sample rate override for custom DS (#117796)

- Add `dynamic-sampling.sample-rate-override-per-project` option — a map
of stringified project id → fixed sample rate — to pin specific projects
to a known rate out-of-band
- Closes the gap where the custom (AM3) DS path had no override hook
(the existing getsentry override only wraps the blended/automatic rate,
which the custom path never calls)
- **Legacy path:** apply the override in
`get_guarded_project_sample_rate`, winning over org/project targets and
the recently-added/rebalanced rate
- **New per-org path:** apply it as an explicit
`apply_project_sample_rate_overrides` step in the scheduler after
balancing, so it's surfaced in the pipeline and feeds the cached rates +
transaction balancing
- Out-of-range/invalid values are ignored (hard-replace semantics
otherwise)
- Tests cover both paths, no-op-without-option, and the AM3
organization-mode baseline behavior

Contributes to TET-2490

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: 2 people

src/sentry/dynamic_sampling/per_org/calculations.py

@@ -2,6 +2,7 @@
 
 import logging
 from collections.abc import Iterable
+from dataclasses import replace
 from typing import TYPE_CHECKING, cast
 
 import orjson
@@ -23,6 +24,7 @@
 from sentry.dynamic_sampling.per_org.gate import project_balancing_debug_project_ids
 from sentry.dynamic_sampling.per_org.queries import ProjectTransactionCounts, ProjectVolume
 from sentry.dynamic_sampling.rules.utils import get_redis_client_for_ds
+from sentry.dynamic_sampling.sample_rate_override import get_sample_rate_overrides
 from sentry.dynamic_sampling.tasks.common import sample_rate_to_float
 from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import (
     generate_boost_low_volume_projects_cache_key,
@@ -63,6 +65,29 @@ def run_project_balancing(
     )
 
 
+def apply_project_sample_rate_overrides(
+    rebalanced_projects: list[RebalancedItem],
+) -> list[RebalancedItem]:
+    """
+    Hard-replace the balanced sample rate of any project that has a per-project override
+    configured via the ``dynamic-sampling.sample-rate-override-per-project`` option.
+
+    Applied as an explicit step in the scheduler (rather than inside the balancing model)
+    so the override is surfaced in the pipeline. The result feeds the cached project
+    sample rates and the downstream transaction balancing.
+    """
+    overrides = get_sample_rate_overrides()
+    if not overrides:
+        return rebalanced_projects
+
+    return [
+        replace(item, new_sample_rate=overrides[int(item.id)])
+        if int(item.id) in overrides
+        else item
+        for item in rebalanced_projects
+    ]
+
+
 def get_cached_rebalanced_project_sample_rates(org_id: int) -> dict[int, float | None]:
     redis_client = get_redis_client_for_ds()
     cache_key = generate_boost_low_volume_projects_cache_key(org_id=org_id)

src/sentry/dynamic_sampling/per_org/scheduler.py

@@ -9,6 +9,7 @@
 from taskbroker_client.retry import Retry
 
 from sentry.dynamic_sampling.per_org.calculations import (
+    apply_project_sample_rate_overrides,
     compare_rebalanced_projects_with_cache,
     compare_rebalanced_transactions_with_cache,
     get_cached_rebalanced_project_sample_rates,
@@ -126,6 +127,7 @@ def run_calculations_per_org_task(org_id: OrganizationId) -> DynamicSamplingStat
 
     if config.should_balance_projects:
         rebalanced_projects = run_project_balancing(config, project_volumes)
+        rebalanced_projects = apply_project_sample_rate_overrides(rebalanced_projects)
         config.set_rebalanced_project_sample_rates(rebalanced_projects)
         cached_sample_rates = get_cached_rebalanced_project_sample_rates(config.organization.id)
         compare_rebalanced_projects_with_cache(

src/sentry/dynamic_sampling/rules/base.py

@@ -9,6 +9,7 @@
 from sentry.dynamic_sampling.rules.biases.base import Bias
 from sentry.dynamic_sampling.rules.combine import get_relay_biases
 from sentry.dynamic_sampling.rules.utils import PolymorphicRule, RuleType, get_enabled_user_biases
+from sentry.dynamic_sampling.sample_rate_override import get_sample_rate_override_for_project
 from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import (
     get_boost_low_volume_projects_sample_rate,
 )
@@ -58,6 +59,12 @@ def get_guarded_project_sample_rate(organization: Organization, project: Project
         return float(project.get_option("sentry:target_sample_rate", TARGET_SAMPLE_RATE_DEFAULT))
 
     if has_custom_dynamic_sampling(organization):
+        # A per-project override (configured via options) hard-replaces the rate the custom
+        # dynamic sampling path would otherwise compute, winning over project/org targets and
+        # the boosted/rebalanced rate.
+        override = get_sample_rate_override_for_project(project.id)
+        if override is not None:
+            return override
         sample_rate = organization.get_option("sentry:target_sample_rate")
     else:
         sample_rate = quotas.backend.get_blended_sample_rate(

src/sentry/dynamic_sampling/sample_rate_override.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from sentry import options
+
+
+def get_sample_rate_overrides() -> dict[int, float]:
+    """
+    Return the validated per-project sample rate overrides for custom dynamic sampling,
+    as configured via the ``dynamic-sampling.sample-rate-override-per-project`` option.
+
+    The option maps a stringified project id to a fixed sample rate that hard-replaces
+    whatever rate the custom dynamic sampling path would otherwise compute. Entries with
+    an invalid id or an out-of-range/invalid rate are skipped (rather than emitting an
+    invalid rule). Reads the option once so callers iterating over many projects don't
+    re-read it per project.
+    """
+    overrides: dict[int, float] = {}
+    for raw_id, raw_rate in options.get(
+        "dynamic-sampling.sample-rate-override-per-project"
+    ).items():
+        try:
+            project_id = int(raw_id)
+            rate = float(raw_rate)
+        except (TypeError, ValueError):
+            continue
+        if 0.0 <= rate <= 1.0:
+            overrides[project_id] = rate
+    return overrides
+
+
+def get_sample_rate_override_for_project(project_id: int) -> float | None:
+    """Return the sample rate override for a single project, or ``None`` if none applies."""
+    return get_sample_rate_overrides().get(project_id)

src/sentry/options/defaults.py

@@ -2066,6 +2066,16 @@
     flags=FLAG_AUTOMATOR_MODIFIABLE,
 )
 
+# Per-project sample rate overrides for custom dynamic sampling. Maps a stringified
+# project id to a fixed sample rate (0.0-1.0) that hard-replaces the rate the custom
+# dynamic sampling path would otherwise compute for that project. Example:
+# {"12345": 0.5}. An empty mapping disables the override.
+register(
+    "dynamic-sampling.sample-rate-override-per-project",
+    default={},
+    flags=FLAG_AUTOMATOR_MODIFIABLE,
+)
+
 # Controls the intensity of dynamic sampling transaction rebalancing. 0.0 = explict rebalancing
 # not performed, 1.0= full rebalancing (tries to bring everything to mean). Note that even at 0.0
 # there will still be some rebalancing between the explicit and implicit transactions ( so setting rebalancing

tests/sentry/dynamic_sampling/per_org/test_calculations.py

@@ -8,6 +8,7 @@
 from sentry.dynamic_sampling.models.common import RebalancedItem
 from sentry.dynamic_sampling.models.projects_rebalancing import ProjectsRebalancingInput
 from sentry.dynamic_sampling.per_org.calculations import (
+    apply_project_sample_rate_overrides,
     compare_rebalanced_projects_with_cache,
     compare_rebalanced_transactions_with_cache,
     get_cached_rebalanced_project_sample_rates,
@@ -25,6 +26,7 @@
     generate_boost_low_volume_transactions_cache_key,
 )
 from sentry.testutils.cases import TestCase
+from sentry.testutils.helpers.options import override_options
 
 
 def _project_volume(project_id: int, total: int = 100, keep: int = 25) -> ProjectVolume:
@@ -71,6 +73,32 @@ def test_run_project_balancing_returns_rebalanced_projects(self) -> None:
         ]
         assert result == rebalanced_projects
 
+    def test_apply_project_sample_rate_overrides(self) -> None:
+        overridden_id = 1001
+        normal_id = 1002
+        rebalanced_projects = [
+            RebalancedItem(id=overridden_id, count=100, new_sample_rate=0.25),
+            RebalancedItem(id=normal_id, count=100, new_sample_rate=0.25),
+        ]
+
+        with override_options(
+            {"dynamic-sampling.sample-rate-override-per-project": {str(overridden_id): 0.9}}
+        ):
+            result = apply_project_sample_rate_overrides(rebalanced_projects)
+
+        result_by_id = {item.id: item.new_sample_rate for item in result}
+        # Overridden project gets the option value; the other keeps its balanced rate.
+        assert result_by_id[overridden_id] == 0.9
+        assert result_by_id[normal_id] == 0.25
+
+    def test_apply_project_sample_rate_overrides_noop_without_option(self) -> None:
+        rebalanced_projects = [
+            RebalancedItem(id=2001, count=100, new_sample_rate=0.25),
+        ]
+        # No overrides configured -> the balanced rates are returned untouched.
+        result = apply_project_sample_rate_overrides(rebalanced_projects)
+        assert result == rebalanced_projects
+
     def test_compare_rebalanced_projects_with_cache_logs_per_project(self) -> None:
         org = self.create_organization()
         project_with_volume = self.create_project(organization=org)

tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py

@@ -27,6 +27,7 @@
 from sentry.testutils.cases import BaseMetricsLayerTestCase, SnubaTestCase, TestCase
 from sentry.testutils.helpers.datetime import freeze_time
 from sentry.testutils.helpers.features import with_feature
+from sentry.testutils.helpers.options import override_options
 
 MOCK_DATETIME = (timezone.now() - timedelta(days=1)).replace(
     hour=0, minute=0, second=0, microsecond=0
@@ -167,6 +168,43 @@ def test_simple_one_org_one_project_task_target_sample_rate(self) -> None:
         )
         assert (sample_rate, got_value) == (0.5, True)
 
+    @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"])
+    def test_per_project_sample_rate_override(self) -> None:
+        # A per-project override configured via options hard-replaces the rate the
+        # custom dynamic sampling path would otherwise resolve for that project --
+        # winning even over the recently-added 100% boost -- and leaves other projects
+        # untouched.
+        org1 = self.create_organization("am3-override-org")
+        org1.update_option("sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION)
+        org1.update_option("sentry:target_sample_rate", 0.5)
+        overridden = self.create_project(organization=org1)
+        normal = self.create_project(organization=org1)
+
+        # Baseline: freshly-created projects are boosted to 1.0 by the recently-added
+        # rule, so neither resolves to the org target yet.
+        assert get_guarded_project_sample_rate(org1, overridden) == 1.0
+
+        with override_options(
+            {"dynamic-sampling.sample-rate-override-per-project": {str(overridden.id): 0.9}}
+        ):
+            assert get_guarded_project_sample_rate(org1, overridden) == 0.9
+            # Not in the override map -> unaffected by the override.
+            assert get_guarded_project_sample_rate(org1, normal) == 1.0
+
+    @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"])
+    def test_per_project_sample_rate_override_ignores_out_of_range(self) -> None:
+        org1 = self.create_organization("am3-override-org-bad")
+        org1.update_option("sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION)
+        org1.update_option("sentry:target_sample_rate", 0.5)
+        project = self.create_project(organization=org1)
+
+        baseline = get_guarded_project_sample_rate(org1, project)
+        with override_options(
+            {"dynamic-sampling.sample-rate-override-per-project": {str(project.id): 2.0}}
+        ):
+            # Out-of-range override is ignored; the resolved rate is unchanged.
+            assert get_guarded_project_sample_rate(org1, project) == baseline
+
     @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"])
     def test_project_mode_sampling_with_query(self) -> None:
         org1 = self.create_organization("test-org")