Repository: getsentry/sentry
Pull request: ref(dynamic-sampling): Implement new abstraction for models
This PR implements a completely new system for defining models that we might use for dynamic sampling.

Closes https://github.com/getsentry/sentry/issues/49930

Problem statement:
Cleanup dynamic sampling code
The dynamic sampling package saw substantial growth w.r.t. responsibilities. Due to the speed of iteration, there is some technical debt in the code which should be cleaned up. It will not require a lot of effort but it will help to keep our standards high.

Hints from issue discussion:

Commit message: Refactor guarded run

Patch to review:
diff --git a/src/sentry/dynamic_sampling/models/adjustment_models.py b/src/sentry/dynamic_sampling/models/adjustment_models.py
deleted file mode 100644
index 39e52303c7b8f1..00000000000000
--- a/src/sentry/dynamic_sampling/models/adjustment_models.py
+++ /dev/null
@@ -1,25 +0,0 @@
-from dataclasses import dataclass
-from typing import List
-
-from sentry.dynamic_sampling.models.utils import DSElement, adjust_sample_rates_full
-
-
-@dataclass
-class AdjustedModel:
-    """
-    Model which can adjust sample rate per project inside ORG based on this new counter metric from Relay.
-    """
-
-    projects: List[DSElement]
-
-    def adjust_sample_rates(self, sample_rate: float) -> List[DSElement]:
-        if len(self.projects) < 2:
-            # When we have one project we just remind sample rates
-            if len(self.projects) == 1:
-                self.projects[0].new_sample_rate = sample_rate
-            return self.projects
-
-        ret_val, _used_budget = adjust_sample_rates_full(
-            self.projects, sample_rate, intensity=1, min_budget=None
-        )
-        return ret_val
diff --git a/src/sentry/dynamic_sampling/models/base.py b/src/sentry/dynamic_sampling/models/base.py
new file mode 100644
index 00000000000000..016261641cdda9
--- /dev/null
+++ b/src/sentry/dynamic_sampling/models/base.py
@@ -0,0 +1,38 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Generic, TypeVar
+
+
+class ModelType(Enum):
+    TRANSACTIONS_REBALANCING = 1
+    PROJECTS_REBALANCING = 2
+    FULL_REBALANCING = 3
+
+
+@dataclass
+class ModelInput(ABC):
+    @abstractmethod
+    def validate(self) -> bool:
+        # By default, we want each model value to be valid.
+        return True
+
+
+class InvalidModelInputError(Exception):
+    pass
+
+
+Input = TypeVar("Input", bound=ModelInput)
+Output = TypeVar("Output")
+
+
+class Model(ABC, Generic[Input, Output]):
+    @abstractmethod
+    def _run(self, model_input: Input) -> Output:
+        raise NotImplementedError()
+
+    def run(self, model_input: Input) -> Output:
+        if not model_input.validate():
+            raise InvalidModelInputError()
+
+        return self._run(model_input)
diff --git a/src/sentry/dynamic_sampling/models/common.py b/src/sentry/dynamic_sampling/models/common.py
new file mode 100644
index 00000000000000..4657aaa4527092
--- /dev/null
+++ b/src/sentry/dynamic_sampling/models/common.py
@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional, Union
+
+import sentry_sdk
+
+from sentry.dynamic_sampling.models.base import Model, ModelInput
+from sentry.dynamic_sampling.rules.utils import ProjectId, TransactionName
+
+
+@dataclass
+class ModelClass:
+    id: Union[ProjectId, TransactionName]
+    count: float
+    new_sample_rate: float = 0.0
+
+
+def sum_classes_counts(classes: List[ModelClass]) -> float:
+    ret_val = 0.0
+
+    for elm in classes:
+        ret_val += elm.count
+
+    return ret_val
+
+
+def guarded_run(model: Model[Any, Any], model_input: ModelInput) -> Optional[Any]:
+    try:
+        return model.run(model_input)
+    except Exception as e:
+        # We want to track the error when running the model.
+        sentry_sdk.capture_exception(e)
+        return None
diff --git a/src/sentry/dynamic_sampling/models/factory.py b/src/sentry/dynamic_sampling/models/factory.py
new file mode 100644
index 00000000000000..445e841dddd0fe
--- /dev/null
+++ b/src/sentry/dynamic_sampling/models/factory.py
@@ -0,0 +1,24 @@
+from typing import TYPE_CHECKING, Any
+
+from sentry.dynamic_sampling.models.base import ModelType
+from sentry.dynamic_sampling.models.full_rebalancing import FullRebalancingModel
+from sentry.dynamic_sampling.models.projects_rebalancing import ProjectsRebalancingModel
+from sentry.dynamic_sampling.models.transactions_rebalancing import TransactionsRebalancingModel
+
+if TYPE_CHECKING:
+    from sentry.dynamic_sampling.models.base import Model
+
+
+class ModelNotFoundError(Exception):
+    pass
+
+
+def model_factory(model_type: ModelType) -> "Model[Any, Any]":
+    if model_type == ModelType.TRANSACTIONS_REBALANCING:
+        return TransactionsRebalancingModel()
+    elif model_type == ModelType.PROJECTS_REBALANCING:
+        return ProjectsRebalancingModel()
+    elif model_type == ModelType.FULL_REBALANCING:
+        return FullRebalancingModel()
+    else:
+        raise ModelNotFoundError()
diff --git a/src/sentry/dynamic_sampling/models/full_rebalancing.py b/src/sentry/dynamic_sampling/models/full_rebalancing.py
new file mode 100644
index 00000000000000..9eafcdf1918f3f
--- /dev/null
+++ b/src/sentry/dynamic_sampling/models/full_rebalancing.py
@@ -0,0 +1,83 @@
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+from sentry.dynamic_sampling.models.base import Model, ModelInput
+from sentry.dynamic_sampling.models.common import ModelClass, sum_classes_counts
+
+
+@dataclass
+class FullRebalancingInput(ModelInput):
+    classes: List[ModelClass]
+    sample_rate: float
+    intensity: float
+    min_budget: Optional[float] = None
+
+    def validate(self) -> bool:
+        return (
+            0.0 <= self.sample_rate <= 1.0
+            and 0.0 <= self.intensity <= 1.0
+            and len(self.classes) > 0
+        )
+
+
+class FullRebalancingModel(Model[FullRebalancingInput, Tuple[List[ModelClass], float]]):
+    def _run(self, model_input: FullRebalancingInput) -> Tuple[List[ModelClass], float]:
+        """
+        Tries to calculate rates that brings all counts close to the ideal count.
+
+        The intensity controls how close, 0 intensity leaves the items unchanged, 1 brings the items to the
+        ideal count ( or rate=1.0 if ideal count is too high).
+
+        :param model_input.classes: The items to be balanced
+        :param model_input.sample_rate: The overall rate necessary
+        :param model_input.intensity: How close to the ideal should we go from our current position (0=do not change, 1 go to ideal)
+        :param model_input.min_budget: Ensure that we use at least min_budget (in order to keep the overall rate)
+
+        :return: A mapping with the frequency for all items + the amount of items used ( it should allways be at least
+        minimum_consumption if passed)
+        """
+        classes = model_input.classes
+        sample_rate = model_input.sample_rate
+        intensity = model_input.intensity
+        min_budget = model_input.min_budget
+
+        total = sum_classes_counts(classes)
+        num_classes = len(classes)
+
+        if min_budget is None:
+            # use exactly what we need (default handling when we resize everything)
+            min_budget = total * sample_rate
+
+        assert total >= min_budget
+        ideal = total * sample_rate / num_classes
+
+        used_budget: float = 0.0
+
+        ret_val = []
+        while classes:
+            element = classes.pop()
+            count = element.count
+            if ideal * num_classes < min_budget:
+                # if we keep to our ideal we will not be able to use the minimum budget (readjust our target)
+                ideal = min_budget / num_classes
+            # see what's the difference from our ideal
+            sampled = count * sample_rate
+            delta = ideal - sampled
+            correction = delta * intensity
+            desired_count = sampled + correction
+
+            if desired_count > count:
+                # we need more than we have, the best we can do is give all, i.e. rate = 1.0
+                new_sample_rate = 1.0
+                used = count
+            else:
+                # we can spend what we want
+                new_sample_rate = desired_count / count
+                used = desired_count
+
+            ret_val.append(ModelClass(id=element.id, count=count, new_sample_rate=new_sample_rate))
+            min_budget -= used
+            used_budget += used
+            num_classes -= 1
+
+        return ret_val, used_budget
diff --git a/src/sentry/dynamic_sampling/models/projects_rebalancing.py b/src/sentry/dynamic_sampling/models/projects_rebalancing.py
new file mode 100644
index 00000000000000..c97c7c86ea9a5e
--- /dev/null
+++ b/src/sentry/dynamic_sampling/models/projects_rebalancing.py
@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+from typing import List
+
+from sentry.dynamic_sampling.models.base import Model, ModelInput, ModelType
+from sentry.dynamic_sampling.models.common import ModelClass
+from sentry.dynamic_sampling.models.full_rebalancing import FullRebalancingInput
+
+
+@dataclass
+class ProjectsRebalancingInput(ModelInput):
+    classes: List[ModelClass]
+    sample_rate: float
+
+    def validate(self) -> bool:
+        return 0.0 <= self.sample_rate <= 1.0
+
+
+class ProjectsRebalancingModel(Model[ProjectsRebalancingInput, List[ModelClass]]):
+    def _run(self, model_input: ProjectsRebalancingInput) -> List[ModelClass]:
+        classes = model_input.classes
+        sample_rate = model_input.sample_rate
+
+        if len(classes) == 0:
+            return classes
+
+        if len(classes) == 1:
+            classes[0].new_sample_rate = sample_rate
+
+        sorted_classes = sorted(classes, key=lambda x: (x.count, x.id), reverse=True)
+
+        from sentry.dynamic_sampling.models.factory import model_factory
+
+        full_rebalancing = model_factory(ModelType.FULL_REBALANCING)
+        result, _ = full_rebalancing.run(
+            FullRebalancingInput(
+                classes=sorted_classes,
+                sample_rate=sample_rate,
+                intensity=1,
+            )
+        )
+
+        return result  # type:ignore
diff --git a/src/sentry/dynamic_sampling/models/transactions_rebalancing.py b/src/sentry/dynamic_sampling/models/transactions_rebalancing.py
new file mode 100644
index 00000000000000..3bf2817235bb54
--- /dev/null
+++ b/src/sentry/dynamic_sampling/models/transactions_rebalancing.py
@@ -0,0 +1,136 @@
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+from sentry.dynamic_sampling.models.base import Model, ModelInput, ModelType
+from sentry.dynamic_sampling.models.common import ModelClass, sum_classes_counts
+from sentry.dynamic_sampling.models.full_rebalancing import FullRebalancingInput
+
+
+@dataclass
+class TransactionsRebalancingInput(ModelInput):
+    classes: List[ModelClass]
+    sample_rate: float
+    total_num_classes: Optional[int]
+    total: Optional[float]
+    intensity: float
+
+    def validate(self) -> bool:
+        return (
+            0.0 <= self.sample_rate <= 1.0
+            and 0.0 <= self.intensity <= 1.0
+            and len(self.classes) > 0
+        )
+
+
+class TransactionsRebalancingModel(
+    Model[TransactionsRebalancingInput, Tuple[List[ModelClass], float]]
+):
+    def _run(self, model_input: TransactionsRebalancingInput) -> Tuple[List[ModelClass], float]:
+        """
+        Adjusts sampling rates to bring the number of samples kept in each class as close to
+        the same value as possible while maintaining the overall sampling rate.
+
+        The algorithm adjusts the explicitly given classes individually to bring them to
+        the ideal sample rate and then adjusts the global sample rate for all the remaining classes.
+
+        :param model_input.classes: a list of class id, num_samples in class
+        :param model_input.sample_rate: global rate of sampling desired
+        :param model_input.total_num_classes: total number of classes (including the explicitly specified in classes)
+        :param model_input.intensity: the adjustment strength 0: no adjustment, 1: try to bring everything to mean
+        :param model_input.total: total number of samples in all classes (including the explicitly specified classes)
+
+        :return: a list of DSElement with calculated sample_rates and  a rate for all other (unspecified) classes.
+        """
+        classes = model_input.classes
+        sample_rate = model_input.sample_rate
+        total_num_classes = model_input.total_num_classes
+        total = model_input.total
+        intensity = model_input.intensity
+
+        classes = sorted(classes, key=lambda x: (x.count, x.id), reverse=True)
+
+        # total count for the explicitly specified classes
+        total_explicit = sum_classes_counts(classes)
+
+        if total is None:
+            total = total_explicit
+
+        if total_num_classes is None:
+            total_num_classes = len(classes)
+
+        # total count for the unspecified classes
+        total_implicit = total - total_explicit
+        # total number of specified classes
+        num_explicit_classes = len(classes)
+        # total number of unspecified classes
+        num_implicit_classes = total_num_classes - num_explicit_classes
+
+        total_budget = total * sample_rate
+        budget_per_class = total_budget / total_num_classes
+
+        implicit_budget = budget_per_class * num_implicit_classes
+        explicit_budget = budget_per_class * num_explicit_classes
+
+        from sentry.dynamic_sampling.models.factory import model_factory
+
+        full_rebalancing = model_factory(ModelType.FULL_REBALANCING)
+
+        if num_explicit_classes == total_num_classes:
+            # we have specified all classes
+            explicit_rates, _used = full_rebalancing.run(
+                FullRebalancingInput(classes=classes, sample_rate=sample_rate, intensity=intensity)
+            )
+            implicit_rate = sample_rate  # doesn't really matter since everything is explicit
+        elif total_implicit < implicit_budget:
+            # we would not be able to spend all implicit budget we can only spend
+            # a maximum of total_implicit, set the implicit rate to 1
+            # and reevaluate the available budget for the explicit classes
+            implicit_rate = 1
+            # we spent all we could on the implicit classes see what budget we
+            # have left
+            explicit_budget = total_budget - total_implicit
+            # calculate the new global rate for the explicit transactions that
+            # would bring the overall rate to the desired rate
+            explicit_rate = explicit_budget / total_explicit
+            explicit_rates, _used = full_rebalancing.run(
+                FullRebalancingInput(
+                    classes=classes, sample_rate=explicit_rate, intensity=intensity
+                )
+            )
+        elif total_explicit < explicit_budget:
+            # we would not be able to spend all explicit budget we can only
+            # send a maximum of total_explicit so set the explicit rate to 1 for
+            # all explicit classes and reevaluate the available budget for the implicit classes
+            explicit_rates = [
+                ModelClass(id=element.id, count=element.count, new_sample_rate=1.0)
+                for element in classes
+            ]
+
+            # calculate the new global rate for the implicit transactions
+            implicit_budget = total_budget - total_explicit
+            implicit_rate = implicit_budget / total_implicit
+        else:
+            # we can spend all the implicit budget on the implicit classes
+            # and all the explicit budget on the explicit classes
+            # see exactly how much we spend on the explicit classes
+            # and leave the rest for the implicit classes
+
+            # calculate what is the minimum amount we need to spend on the
+            # explicit classes (so that we maintain the overall rate)
+            # if it is <= 0 then we don't have a minimum
+            minimum_explicit_budget = total_budget - total_implicit
+            explicit_rate = explicit_budget / total_explicit
+
+            explicit_rates, used = full_rebalancing.run(
+                FullRebalancingInput(
+                    classes=classes,
+                    sample_rate=explicit_rate,
+                    intensity=intensity,
+                    min_budget=minimum_explicit_budget,
+                )
+            )
+            # recalculate implicit_budget based on used
+            implicit_budget = total_budget - used
+            implicit_rate = implicit_budget / total_implicit
+
+        return explicit_rates, implicit_rate
diff --git a/src/sentry/dynamic_sampling/models/utils.py b/src/sentry/dynamic_sampling/models/utils.py
deleted file mode 100644
index 3077ec938d4005..00000000000000
--- a/src/sentry/dynamic_sampling/models/utils.py
+++ /dev/null
@@ -1,187 +0,0 @@
-from dataclasses import dataclass
-from typing import List, Optional, Sequence, Tuple, Union
-
-from sentry.dynamic_sampling.rules.utils import ProjectId, TransactionName
-
-
-@dataclass
-class DSElement:
-    id: Union[ProjectId, TransactionName]
-    count: float
-    new_sample_rate: float = 0.0
-
-
-def adjust_sample_rates(
-    classes: List[DSElement],
-    rate: float,
-    total_num_classes: Optional[int],
-    total: Optional[float],
-    intensity: float,
-) -> Tuple[List[DSElement], float]:
-    """
-    Adjusts sampling rates to bring the number of samples kept in each class as close to
-    the same value as possible while maintaining the overall sampling rate.
-
-    The algorithm adjusts the explicitly given classes individually to bring them to
-    the ideal sample rate and then adjusts the global sample rate for all the remaining classes.
-
-    :param classes: a list of class id, num_samples in class
-    :param rate: global rate of sampling desired
-    :param total_num_classes: total number of classes (including the explicitly specified in classes)
-    :param intensity: the adjustment strength 0: no adjustment, 1: try to bring everything to mean
-    :param total: total number of samples in all classes (including the explicitly specified classes)
-
-    :return: a list of DSElement with calculated sample_rates and  a rate for all other (unspecified) classes.
-    """
-
-    classes = sorted(classes, key=lambda x: (x.count, x.id), reverse=True)
-
-    # total count for the explicitly specified classes
-    total_explicit = sum_counts(classes)
-
-    if total is None:
-        total = total_explicit
-
-    if total_num_classes is None:
-        total_num_classes = len(classes)
-
-    # total count for the unspecified classes
-    total_implicit = total - total_explicit
-    # total number of specified classes
-    num_explicit_classes = len(classes)
-    # total number of unspecified classes
-    num_implicit_classes = total_num_classes - num_explicit_classes
-
-    total_budget = total * rate
-    budget_per_class = total_budget / total_num_classes
-
-    implicit_budget = budget_per_class * num_implicit_classes
-    explicit_budget = budget_per_class * num_explicit_classes
-
-    if num_explicit_classes == total_num_classes:
-        # we have specified all classes
-        explicit_rates, _used = _adjust_sample_rates_full(
-            classes, rate=rate, intensity=intensity, min_budget=None
-        )
-        implicit_rate = rate  # doesn't really matter since everything is explicit
-    elif total_implicit < implicit_budget:
-        # we would not be able to spend all implicit budget we can only spend
-        # a maximum of total_implicit, set the implicit rate to 1
-        # and reevaluate the available budget for the explicit classes
-        implicit_rate = 1
-        # we spent all we could on the implicit classes see what budget we
-        # have left
-        explicit_budget = total_budget - total_implicit
-        # calculate the new global rate for the explicit transactions that
-        # would bring the overall rate to the desired rate
-        explicit_rate = explicit_budget / total_explicit
-        explicit_rates, _used = _adjust_sample_rates_full(
-            classes, explicit_rate, intensity=intensity, min_budget=None
-        )
-    elif total_explicit < explicit_budget:
-        # we would not be able to spend all explicit budget we can only
-        # send a maximum of total_explicit so set the explicit rate to 1 for
-        # all explicit classes and reevaluate the available budget for the implicit classes
-        explicit_rates = [
-            DSElement(id=element.id, count=element.count, new_sample_rate=1.0)
-            for element in classes
-        ]
-
-        # calculate the new global rate for the implicit transactions
-        implicit_budget = total_budget - total_explicit
-        implicit_rate = implicit_budget / total_implicit
-    else:
-        # we can spend all the implicit budget on the implicit classes
-        # and all the explicit budget on the explicit classes
-        # see exactly how much we spend on the explicit classes
-        # and leave the rest for the implicit classes
-
-        # calculate what is the minimum amount we need to spend on the
-        # explicit classes (so that we maintain the overall rate)
-        # if it is <= 0 then we don't have a minimum
-        minimum_explicit_budget = total_budget - total_implicit
-        explicit_rate = explicit_budget / total_explicit
-
-        explicit_rates, used = _adjust_sample_rates_full(
-            classes=classes,
-            rate=explicit_rate,
-            intensity=intensity,
-            min_budget=minimum_explicit_budget,
-        )
-        # recalculate implicit_budget based on used
-        implicit_budget = total_budget - used
-        implicit_rate = implicit_budget / total_implicit
-    return explicit_rates, implicit_rate
-
-
-def adjust_sample_rates_full(
-    classes: Sequence[DSElement], rate: float, intensity: float, min_budget: Optional[float] = None
-) -> Tuple[List[DSElement], float]:
-    sorted_elements = sorted(classes, key=lambda x: (x.count, x.id), reverse=True)
-    return _adjust_sample_rates_full(sorted_elements, rate, intensity, min_budget)
-
-
-def _adjust_sample_rates_full(
-    classes: List[DSElement], rate: float, intensity: float, min_budget: Optional[float] = None
-) -> Tuple[List[DSElement], float]:
-    """
-    Tries to calculate rates that brings all counts close to the ideal count
-
-    The intensity controls how close, 0 intensity leaves the items unchanged, 1 brings the items to the
-    ideal count ( or rate=1.0 if ideal count is too high).
-
-    :param items: The items to be balanced
-    :param rate: The overall rate necessary
-    :param intensity: How close to the ideal should we go from our current position (0=do not change, 1 go to ideal)
-    :param min_budget: Ensure that we use at least min_budget (in order to keep the overall rate)
-
-    :return: A mapping with the frequency for all items + the amount of items used ( it should allways be at least
-    minimum_consumption if passed)
-    """
-
-    total = sum_counts(classes)
-    num_classes = len(classes)
-
-    if min_budget is None:
-        # use exactly what we need (default handling when we resize everything)
-        min_budget = total * rate
-
-    assert total >= min_budget
-    ideal = total * rate / num_classes
-
-    used_budget: float = 0.0
-
-    ret_val = []
-    while classes:
-        element = classes.pop()
-        count = element.count
-        if ideal * num_classes < min_budget:
-            # if we keep to our ideal we will not be able to use the minimum budget (readjust our target)
-            ideal = min_budget / num_classes
-        # see what's the difference from our ideal
-        sampled = count * rate
-        delta = ideal - sampled
-        correction = delta * intensity
-        desired_count = sampled + correction
-
-        if desired_count > count:
-            # we need more than we have, the best we can do is give all, i.e. rate = 1.0
-            new_sample_rate = 1.0
-            used = count
-        else:
-            # we can spend what we want
-            new_sample_rate = desired_count / count
-            used = desired_count
-        ret_val.append(DSElement(id=element.id, count=count, new_sample_rate=new_sample_rate))
-        min_budget -= used
-        used_budget += used
-        num_classes -= 1
-
-    return ret_val, used_budget
-
-
-def sum_counts(transactions: List[DSElement]) -> float:
-    ret_val = 0.0
-    for elm in transactions:
-        ret_val += elm.count
-    return ret_val
diff --git a/src/sentry/dynamic_sampling/rules/helpers/prioritize_transactions.py b/src/sentry/dynamic_sampling/rules/helpers/prioritize_transactions.py
index 837a20e243a107..e3b61fd6280910 100644
--- a/src/sentry/dynamic_sampling/rules/helpers/prioritize_transactions.py
+++ b/src/sentry/dynamic_sampling/rules/helpers/prioritize_transactions.py
@@ -2,7 +2,7 @@
 
 import sentry_sdk
 
-from sentry.dynamic_sampling.models.utils import DSElement
+from sentry.dynamic_sampling.models.common import ModelClass
 from sentry.dynamic_sampling.rules.utils import get_redis_client_for_ds
 from sentry.utils import json
 
@@ -29,7 +29,7 @@ def get_transactions_resampling_rates(
 
 
 def set_transactions_resampling_rates(
-    org_id: int, proj_id: int, named_rates: List[DSElement], default_rate: float, ttl_ms: int
+    org_id: int, proj_id: int, named_rates: List[ModelClass], default_rate: float, ttl_ms: int
 ) -> None:
     redis_client = get_redis_client_for_ds()
     cache_key = _get_cache_key(org_id=org_id, proj_id=proj_id)
diff --git a/src/sentry/dynamic_sampling/tasks.py b/src/sentry/dynamic_sampling/tasks.py
index e24c7c4281e127..b954e86d7045af 100644
--- a/src/sentry/dynamic_sampling/tasks.py
+++ b/src/sentry/dynamic_sampling/tasks.py
@@ -6,9 +6,11 @@
 import sentry_sdk
 
 from sentry import options, quotas
-from sentry.dynamic_sampling.models import utils
-from sentry.dynamic_sampling.models.adjustment_models import AdjustedModel
-from sentry.dynamic_sampling.models.utils import DSElement
+from sentry.dynamic_sampling.models.base import ModelType
+from sentry.dynamic_sampling.models.common import ModelClass, guarded_run
+from sentry.dynamic_sampling.models.factory import model_factory
+from sentry.dynamic_sampling.models.projects_rebalancing import ProjectsRebalancingInput
+from sentry.dynamic_sampling.models.transactions_rebalancing import TransactionsRebalancingInput
 from sentry.dynamic_sampling.prioritise_projects import fetch_projects_with_total_volumes
 from sentry.dynamic_sampling.prioritise_transactions import (
     ProjectTransactions,
@@ -207,7 +209,7 @@ def process_transaction_biases(project_transactions: ProjectTransactions) -> Non
     total_num_transactions = project_transactions.get("total_num_transactions")
     total_num_classes = project_transactions.get("total_num_classes")
     transactions = [
-        DSElement(id=id, count=count) for id, count in project_transactions["transaction_counts"]
+        ModelClass(id=id, count=count) for id, count in project_transactions["transaction_counts"]
     ]
 
     try:
@@ -244,14 +246,24 @@ def process_transaction_biases(project_transactions: ProjectTransactions) -> Non
         return
 
     intensity = options.get("dynamic-sampling.prioritise_transactions.rebalance_intensity", 1.0)
-    named_rates, implicit_rate = utils.adjust_sample_rates(
-        classes=transactions,
-        rate=sample_rate,
-        total_num_classes=total_num_classes,
-        total=total_num_transactions,
-        intensity=intensity,
+
+    model = model_factory(ModelType.TRANSACTIONS_REBALANCING)
+    rebalanced_transactions = guarded_run(
+        model,
+        TransactionsRebalancingInput(
+            classes=transactions,
+            sample_rate=sample_rate,
+            total_num_classes=total_num_classes,
+            total=total_num_transactions,
+            intensity=intensity,
+        ),
     )
+    # In case the result of the model is None, it means that an error occurred, thus we want to early return.
+    if rebalanced_transactions is None:
+        return
 
+    # Only after checking the nullability of rebalanced_transactions, we want to unpack the tuple.
+    named_rates, implicit_rate = rebalanced_transactions
     set_transactions_resampling_rates(
         org_id=org_id,
         proj_id=project_id,
@@ -360,35 +372,43 @@ def adjust_sample_rates(
     projects = []
     for project_id, count_per_root in projects_with_counts.items():
         projects.append(
-            DSElement(
+            ModelClass(
                 id=project_id,
                 count=count_per_root,
             )
         )
 
-    model = AdjustedModel(projects=projects)
-    ds_projects = model.adjust_sample_rates(sample_rate=sample_rate)
+    model = model_factory(ModelType.PROJECTS_REBALANCING)
+    rebalanced_projects = guarded_run(
+        model, ProjectsRebalancingInput(classes=projects, sample_rate=sample_rate)
+    )
+    # In case the result of the model is None, it means that an error occurred, thus we want to early return.
+    if rebalanced_projects is None:
+        return
 
     redis_client = get_redis_client_for_ds()
     with redis_client.pipeline(transaction=False) as pipeline:
-        for ds_project in ds_projects:
+        for rebalanced_project in rebalanced_projects:
             cache_key = generate_prioritise_by_project_cache_key(org_id=org_id)
             # We want to get the old sample rate, which will be None in case it was not set.
-            old_sample_rate = sample_rate_to_float(redis_client.hget(cache_key, ds_project.id))
+            old_sample_rate = sample_rate_to_float(
+                redis_client.hget(cache_key, rebalanced_project.id)
+            )
 
             # We want to store the new sample rate as a string.
             pipeline.hset(
                 cache_key,
-                ds_project.id,
-                ds_project.new_sample_rate,  # redis stores is as string
+                rebalanced_project.id,
+                rebalanced_project.new_sample_rate,  # redis stores is as string
             )
             pipeline.pexpire(cache_key, CACHE_KEY_TTL)
 
             # We invalidate the caches only if there was a change in the sample rate. This is to avoid flooding the
             # system with project config invalidations, especially for projects with no volume.
-            if not are_equal_with_epsilon(old_sample_rate, ds_project.new_sample_rate):
+            if not are_equal_with_epsilon(old_sample_rate, rebalanced_project.new_sample_rate):
                 schedule_invalidate_project_config(
-                    project_id=ds_project.id, trigger="dynamic_sampling_prioritise_project_bias"
+                    project_id=rebalanced_project.id,
+                    trigger="dynamic_sampling_prioritise_project_bias",
                 )
 
         pipeline.execute()
diff --git a/tests/sentry/dynamic_sampling/models/test_adjusments_models.py b/tests/sentry/dynamic_sampling/models/test_adjusments_models.py
deleted file mode 100644
index 456984862b9ea4..00000000000000
--- a/tests/sentry/dynamic_sampling/models/test_adjusments_models.py
+++ /dev/null
@@ -1,220 +0,0 @@
-from operator import attrgetter
-
-import pytest
-
-from sentry.dynamic_sampling.models.adjustment_models import AdjustedModel
-from sentry.dynamic_sampling.models.adjustment_models import DSElement as P
-
-
-def test_adjust_sample_rates_org_with_no_projects():
-    p = AdjustedModel(projects=[])
-    assert p.adjust_sample_rates(sample_rate=0.25) == []
-
-
-def test_adjust_sample_rates_org_with_single_project():
-    p = AdjustedModel(
-        projects=[
-            P(
-                id=1,
-                count=10,
-            )
-        ]
-    )
-    assert p.adjust_sample_rates(sample_rate=0.4) == [P(id=1, count=10, new_sample_rate=0.4)]
-
-
-def test_adjust_sample_rates_org_with_few_projects():
-    projects = [
-        P(id=1, count=9),
-        P(id=2, count=7),
-        P(id=3, count=3),
-        P(id=4, count=1),
-    ]
-    p = AdjustedModel(projects=projects)
-
-    expected_projects = [
-        P(
-            id=1,
-            count=9,
-            new_sample_rate=pytest.approx(0.14814814814814817),
-        ),
-        P(
-            id=2,
-            count=7,
-            new_sample_rate=pytest.approx(0.1904761904761905),
-        ),
-        P(
-            id=3,
-            count=3,
-            new_sample_rate=pytest.approx(0.4444444444444444),
-        ),
-        P(
-            id=4,
-            count=1,
-            new_sample_rate=1.0,
-        ),
-    ]
-
-    assert (
-        sorted(p.adjust_sample_rates(sample_rate=0.25), key=attrgetter("id")) == expected_projects
-    )
-
-
-def test_adjust_sample_rates_org_with_many_projects():
-    """
-    This test checks how we calculate sample rates for org with 30 projects
-    and make sure model doesn't generate negative sample rate
-    """
-    target = [
-        (1, 13369016.0),
-        (11276, 369985.0),
-        (37617, 55.0),
-        (155735, 126.0),
-        (162676, 10550.0),
-        (191772, 5.0),
-        (239368, 1.0),
-        (300688, 2595111.0),
-        (1267915, 15199.0),
-        (1886021, 445.0),
-        (2053674, 47.0),
-        (4857230, 26530.0),
-        (5246761, 8211.0),
-        (5266138, 737.0),
-        (5324467, 89.0),
-        (5350637, 11.0),
-        (5600888, 68.0),
-        (5613292, 161.0),
-        (5683166, 63.0),
-        (5738630, 257.0),
-        (5899451, 1417.0),
-        (5903949, 10263.0),
-        (6178942, 358470.0),
-        (6301746, 244.0),
-        (6418660, 21.0),
-        (6424467, 5055781.0),
-        (6690737, 773.0),
-        (4504044639748096, 493.0),
-        (4504044642107392, 1564.0),
-        (4504373448540160, 22.0),
-    ]
-    projects = [
-        P(
-            id=p_id,
-            count=count,
-        )
-        for p_id, count in target
-    ]
-    p = AdjustedModel(projects=projects)
-
-    sample_rate = 0.036
-    result = p.adjust_sample_rates(sample_rate)
-    for p in result:
-        assert 1.0 >= p.new_sample_rate > 0
-        assert (p.count * sample_rate) / (
-            p.count * p.new_sample_rate
-        ) * p.new_sample_rate == pytest.approx(sample_rate)
-
-
-def test_adjust_sample_rates_org_with_even_num_projects():
-    projects = [
-        P(id=1, count=8.0),
-        P(id=2, count=7.0),
-        P(id=3, count=3.0),
-    ]
-    p = AdjustedModel(projects=projects)
-
-    expected_projects = [
-        P(
-            id=1,
-            count=8.0,
-            new_sample_rate=0.1875,
-        ),
-        P(
-            id=2,
-            count=7.0,
-            new_sample_rate=pytest.approx(0.21428571428571427),
-        ),
-        P(
-            id=3,
-            count=3.0,
-            new_sample_rate=0.5,
-        ),
-    ]
-
-    assert (
-        sorted(p.adjust_sample_rates(sample_rate=0.25), key=attrgetter("id")) == expected_projects
-    )
-
-
-def test_adjust_sample_rates_org_with_same_counts_projects():
-    projects = [
-        P(id=1, count=9.0),
-        P(id=2, count=6.0),
-        P(id=3, count=6.0),
-        P(id=4, count=1.0),
-    ]
-    p = AdjustedModel(projects=projects)
-
-    expected_projects = [
-        P(
-            id=1,
-            count=9.0,
-            new_sample_rate=pytest.approx(0.16666666666666666),
-        ),
-        P(
-            id=2,
-            count=6.0,
-            new_sample_rate=0.25,
-        ),
-        P(
-            id=3,
-            count=6.0,
-            new_sample_rate=0.25,
-        ),
-        P(
-            id=4,
-            count=1.0,
-            new_sample_rate=1.0,
-        ),
-    ]
-
-    assert (
-        sorted(p.adjust_sample_rates(sample_rate=0.25), key=attrgetter("id")) == expected_projects
-    )
-
-
-def test_adjust_sample_rates_org_with_counts_projects():
-    projects = [
-        P(id=1, count=2.0),
-        P(id=2, count=10.0),
-        P(id=3, count=10.0),
-        P(id=4, count=10.0),
-    ]
-    p = AdjustedModel(projects=projects)
-
-    expected_projects = [
-        P(
-            id=1,
-            count=2.0,
-            new_sample_rate=1.0,
-        ),
-        P(
-            id=2,
-            count=10.0,
-            new_sample_rate=0.2,
-        ),
-        P(
-            id=3,
-            count=10.0,
-            new_sample_rate=0.2,
-        ),
-        P(
-            id=4,
-            count=10.0,
-            new_sample_rate=0.2,
-        ),
-    ]
-
-    assert (
-        sorted(p.adjust_sample_rates(sample_rate=0.25), key=attrgetter("id")) == expected_projects
-    )
diff --git a/tests/sentry/dynamic_sampling/models/test_base.py b/tests/sentry/dynamic_sampling/models/test_base.py
new file mode 100644
index 00000000000000..9e633055c15dde
--- /dev/null
+++ b/tests/sentry/dynamic_sampling/models/test_base.py
@@ -0,0 +1,34 @@
+import pytest
+
+from sentry.dynamic_sampling.models.base import InvalidModelInputError, ModelType
+from sentry.dynamic_sampling.models.factory import model_factory
+from sentry.dynamic_sampling.models.full_rebalancing import FullRebalancingInput
+
+
+@pytest.fixture
+def full_rebalancing_model():
+    return model_factory(ModelType.FULL_REBALANCING)
+
+
+def test_run_with_exception(full_rebalancing_model):
+    with pytest.raises(InvalidModelInputError):
+        full_rebalancing_model.run(
+            FullRebalancingInput(
+                classes=[],
+                sample_rate=0.0,
+                intensity=1,
+            )
+        )
+
+
+def test_guarded_run_with_exception(full_rebalancing_model):
+    assert (
+        full_rebalancing_model.guarded_run(
+            FullRebalancingInput(
+                classes=[],
+                sample_rate=0.0,
+                intensity=1,
+            )
+        )
+        is None
+    )
diff --git a/tests/sentry/dynamic_sampling/models/test_projects_rebalancing.py b/tests/sentry/dynamic_sampling/models/test_projects_rebalancing.py
new file mode 100644
index 00000000000000..ddf7f56a97a498
--- /dev/null
+++ b/tests/sentry/dynamic_sampling/models/test_projects_rebalancing.py
@@ -0,0 +1,252 @@
+from operator import attrgetter
+
+import pytest
+
+from sentry.dynamic_sampling.models.base import ModelType
+from sentry.dynamic_sampling.models.common import ModelClass
+from sentry.dynamic_sampling.models.factory import model_factory
+from sentry.dynamic_sampling.models.projects_rebalancing import ProjectsRebalancingInput
+
+
+@pytest.fixture
+def projects_rebalancing_model():
+    return model_factory(ModelType.PROJECTS_REBALANCING)
+
+
+def test_adjust_sample_rates_org_with_no_projects(projects_rebalancing_model):
+    assert (
+        projects_rebalancing_model.run(ProjectsRebalancingInput(classes=[], sample_rate=0.25)) == []
+    )
+
+
+def test_adjust_sample_rates_org_with_single_project(projects_rebalancing_model):
+    assert projects_rebalancing_model.run(
+        ProjectsRebalancingInput(
+            classes=[
+                ModelClass(
+                    id=1,
+                    count=10,
+                )
+            ],
+            sample_rate=0.4,
+        )
+    ) == [ModelClass(id=1, count=10, new_sample_rate=0.4)]
+
+
+def test_adjust_sample_rates_org_with_few_projects(projects_rebalancing_model):
+    classes = [
+        ModelClass(id=1, count=9),
+        ModelClass(id=2, count=7),
+        ModelClass(id=3, count=3),
+        ModelClass(id=4, count=1),
+    ]
+
+    expected_classes = [
+        ModelClass(
+            id=1,
+            count=9,
+            new_sample_rate=pytest.approx(0.14814814814814817),
+        ),
+        ModelClass(
+            id=2,
+            count=7,
+            new_sample_rate=pytest.approx(0.1904761904761905),
+        ),
+        ModelClass(
+            id=3,
+            count=3,
+            new_sample_rate=pytest.approx(0.4444444444444444),
+        ),
+        ModelClass(
+            id=4,
+            count=1,
+            new_sample_rate=1.0,
+        ),
+    ]
+
+    assert (
+        sorted(
+            projects_rebalancing_model.run(
+                ProjectsRebalancingInput(classes=classes, sample_rate=0.25)
+            ),
+            key=attrgetter("id"),
+        )
+        == expected_classes
+    )
+
+
+def test_adjust_sample_rates_org_with_many_projects(projects_rebalancing_model):
+    """
+    This test checks how we calculate sample rates for org with 30 projects
+    and make sure model doesn't generate negative sample rate
+    """
+    target = [
+        (1, 13369016.0),
+        (11276, 369985.0),
+        (37617, 55.0),
+        (155735, 126.0),
+        (162676, 10550.0),
+        (191772, 5.0),
+        (239368, 1.0),
+        (300688, 2595111.0),
+        (1267915, 15199.0),
+        (1886021, 445.0),
+        (2053674, 47.0),
+        (4857230, 26530.0),
+        (5246761, 8211.0),
+        (5266138, 737.0),
+        (5324467, 89.0),
+        (5350637, 11.0),
+        (5600888, 68.0),
+        (5613292, 161.0),
+        (5683166, 63.0),
+        (5738630, 257.0),
+        (5899451, 1417.0),
+        (5903949, 10263.0),
+        (6178942, 358470.0),
+        (6301746, 244.0),
+        (6418660, 21.0),
+        (6424467, 5055781.0),
+        (6690737, 773.0),
+        (4504044639748096, 493.0),
+        (4504044642107392, 1564.0),
+        (4504373448540160, 22.0),
+    ]
+
+    classes = [
+        ModelClass(
+            id=p_id,
+            count=count,
+        )
+        for p_id, count in target
+    ]
+
+    sample_rate = 0.036
+    result = projects_rebalancing_model.run(
+        ProjectsRebalancingInput(classes=classes, sample_rate=sample_rate)
+    )
+    for p in result:
+        assert 1.0 >= p.new_sample_rate > 0
+        assert (p.count * sample_rate) / (
+            p.count * p.new_sample_rate
+        ) * p.new_sample_rate == pytest.approx(sample_rate)
+
+
+def test_adjust_sample_rates_org_with_even_num_projects(projects_rebalancing_model):
+    classes = [
+        ModelClass(id=1, count=8.0),
+        ModelClass(id=2, count=7.0),
+        ModelClass(id=3, count=3.0),
+    ]
+
+    expected_classes = [
+        ModelClass(
+            id=1,
+            count=8.0,
+            new_sample_rate=0.1875,
+        ),
+        ModelClass(
+            id=2,
+            count=7.0,
+            new_sample_rate=pytest.approx(0.21428571428571427),
+        ),
+        ModelClass(
+            id=3,
+            count=3.0,
+            new_sample_rate=0.5,
+        ),
+    ]
+
+    assert (
+        sorted(
+            projects_rebalancing_model.run(
+                ProjectsRebalancingInput(classes=classes, sample_rate=0.25)
+            ),
+            key=attrgetter("id"),
+        )
+        == expected_classes
+    )
+
+
+def test_adjust_sample_rates_org_with_same_counts_projects(projects_rebalancing_model):
+    classes = [
+        ModelClass(id=1, count=9.0),
+        ModelClass(id=2, count=6.0),
+        ModelClass(id=3, count=6.0),
+        ModelClass(id=4, count=1.0),
+    ]
+
+    expected_classes = [
+        ModelClass(
+            id=1,
+            count=9.0,
+            new_sample_rate=pytest.approx(0.16666666666666666),
+        ),
+        ModelClass(
+            id=2,
+            count=6.0,
+            new_sample_rate=0.25,
+        ),
+        ModelClass(
+            id=3,
+            count=6.0,
+            new_sample_rate=0.25,
+        ),
+        ModelClass(
+            id=4,
+            count=1.0,
+            new_sample_rate=1.0,
+        ),
+    ]
+
+    assert (
+        sorted(
+            projects_rebalancing_model.run(
+                ProjectsRebalancingInput(classes=classes, sample_rate=0.25)
+            ),
+            key=attrgetter("id"),
+        )
+        == expected_classes
+    )
+
+
+def test_adjust_sample_rates_org_with_counts_projects(projects_rebalancing_model):
+    classes = [
+        ModelClass(id=1, count=2.0),
+        ModelClass(id=2, count=10.0),
+        ModelClass(id=3, count=10.0),
+        ModelClass(id=4, count=10.0),
+    ]
+
+    expected_classes = [
+        ModelClass(
+            id=1,
+            count=2.0,
+            new_sample_rate=1.0,
+        ),
+        ModelClass(
+            id=2,
+            count=10.0,
+            new_sample_rate=0.2,
+        ),
+        ModelClass(
+            id=3,
+            count=10.0,
+            new_sample_rate=0.2,
+        ),
+        ModelClass(
+            id=4,
+            count=10.0,
+            new_sample_rate=0.2,
+        ),
+    ]
+
+    assert (
+        sorted(
+            projects_rebalancing_model.run(
+                ProjectsRebalancingInput(classes=classes, sample_rate=0.25)
+            ),
+            key=attrgetter("id"),
+        )
+        == expected_classes
+    )
diff --git a/tests/sentry/dynamic_sampling/models/test_transaction_adjustment_model.py b/tests/sentry/dynamic_sampling/models/test_transactions_rebalancing.py
similarity index 66%
rename from tests/sentry/dynamic_sampling/models/test_transaction_adjustment_model.py
rename to tests/sentry/dynamic_sampling/models/test_transactions_rebalancing.py
index 8b6b7f72262f75..543920b5ae8e5b 100644
--- a/tests/sentry/dynamic_sampling/models/test_transaction_adjustment_model.py
+++ b/tests/sentry/dynamic_sampling/models/test_transactions_rebalancing.py
@@ -2,13 +2,21 @@
 
 import pytest
 
-from sentry.dynamic_sampling.models.utils import DSElement, adjust_sample_rates, sum_counts
+from sentry.dynamic_sampling.models.base import ModelType
+from sentry.dynamic_sampling.models.common import ModelClass, sum_classes_counts
+from sentry.dynamic_sampling.models.factory import model_factory
+from sentry.dynamic_sampling.models.transactions_rebalancing import TransactionsRebalancingInput
+
+
+@pytest.fixture
+def transactions_rebalancing_model():
+    return model_factory(ModelType.TRANSACTIONS_REBALANCING)
 
 
 def create_transaction_counts(big: int, med: int, small: int):
-    big_t = [DSElement(id=f"tb{i}", count=1000 + i) for i in range(big)]
-    med_t = [DSElement(id=f"tm{i}", count=100 + i) for i in range(med)]
-    small_t = [DSElement(id=f"ts{i}", count=1 + i) for i in range(small)]
+    big_t = [ModelClass(id=f"tb{i}", count=1000 + i) for i in range(big)]
+    med_t = [ModelClass(id=f"tm{i}", count=100 + i) for i in range(med)]
+    small_t = [ModelClass(id=f"ts{i}", count=1 + i) for i in range(small)]
     return [*big_t, *med_t, *small_t]
 
 
@@ -32,7 +40,7 @@ def create_transaction_counts(big: int, med: int, small: int):
 
 
 def get_num_sampled_elements(
-    transactions: List[DSElement], trans_dict: Mapping[str, float], global_rate: float
+    transactions: List[ModelClass], trans_dict: Mapping[str, float], global_rate: float
 ) -> float:
     num_transactions = 0.0
     for transaction in transactions:
@@ -48,20 +56,24 @@ def get_num_sampled_elements(
 @pytest.mark.parametrize("sample_rate", sample_rates)
 @pytest.mark.parametrize("transactions", test_resample_cases)
 @pytest.mark.parametrize("idx_low,idx_high", excluded_transactions)
-def test_maintains_overall_sample_rate(intensity, sample_rate, transactions, idx_low, idx_high):
+def test_maintains_overall_sample_rate(
+    transactions_rebalancing_model, intensity, sample_rate, transactions, idx_low, idx_high
+):
     """
     Tests that the overall sampling rate is maintained after applying new rates
     """
     explict_transactions = transactions[idx_low:idx_high]
-    total = sum_counts(transactions)
+    total = sum_classes_counts(transactions)
     total_classes = len(transactions)
 
-    trans, global_rate = adjust_sample_rates(
-        explict_transactions,
-        sample_rate,
-        total_num_classes=total_classes,
-        total=total,
-        intensity=1,
+    trans, global_rate = transactions_rebalancing_model.run(
+        TransactionsRebalancingInput(
+            classes=explict_transactions,
+            sample_rate=sample_rate,
+            total_num_classes=total_classes,
+            total=total,
+            intensity=1,
+        ),
     )
 
     trans_dict = {t.id: t.new_sample_rate for t in trans}
@@ -76,7 +88,9 @@ def test_maintains_overall_sample_rate(intensity, sample_rate, transactions, idx
 @pytest.mark.parametrize("sample_rate", sample_rates)
 @pytest.mark.parametrize("transactions", test_resample_cases)
 @pytest.mark.parametrize("idx_low,idx_high", excluded_transactions)
-def test_explicit_elements_ideal_rate(sample_rate, transactions, idx_low, idx_high):
+def test_explicit_elements_ideal_rate(
+    transactions_rebalancing_model, sample_rate, transactions, idx_low, idx_high
+):
     """
     Tests that the explicitly specified elements are sampled at their ideal rate.
 
@@ -86,15 +100,17 @@ def test_explicit_elements_ideal_rate(sample_rate, transactions, idx_low, idx_hi
     * the budget per transaction
     """
     explict_transactions = transactions[idx_low:idx_high]
-    total = sum_counts(transactions)
+    total = sum_classes_counts(transactions)
     total_classes = len(transactions)
 
-    trans, global_rate = adjust_sample_rates(
-        explict_transactions,
-        sample_rate,
-        total_num_classes=total_classes,
-        total=total,
-        intensity=1.0,
+    trans, global_rate = transactions_rebalancing_model.run(
+        TransactionsRebalancingInput(
+            classes=explict_transactions,
+            sample_rate=sample_rate,
+            total_num_classes=total_classes,
+            total=total,
+            intensity=1,
+        ),
     )
 
     ideal_number_of_elements_per_class = total * sample_rate / total_classes
diff --git a/tests/sentry/dynamic_sampling/rules/helpers/test_prioritize_transactions.py b/tests/sentry/dynamic_sampling/rules/helpers/test_prioritize_transactions.py
index c0ac22f14aa5de..a52fe7b0640a5b 100644
--- a/tests/sentry/dynamic_sampling/rules/helpers/test_prioritize_transactions.py
+++ b/tests/sentry/dynamic_sampling/rules/helpers/test_prioritize_transactions.py
@@ -1,4 +1,4 @@
-from sentry.dynamic_sampling.models.utils import DSElement
+from sentry.dynamic_sampling.models.common import ModelClass
 from sentry.dynamic_sampling.rules.helpers.prioritize_transactions import (
     get_transactions_resampling_rates,
     set_transactions_resampling_rates,
@@ -17,13 +17,13 @@ def test_resampling_rates_in_cache():
 
     expected_global_rate = 0.3
     expected_trans_rates = [
-        DSElement(id="t1", count=1, new_sample_rate=0.6),
-        DSElement(id="t2", count=1, new_sample_rate=0.7),
+        ModelClass(id="t1", count=1, new_sample_rate=0.6),
+        ModelClass(id="t2", count=1, new_sample_rate=0.7),
     ]
     expected_trans_rates_dict = {elm.id: elm.new_sample_rate for elm in expected_trans_rates}
 
     other_rate = 0.1
-    other_rates = [DSElement(id="t11", new_sample_rate=0.1, count=1)]
+    other_rates = [ModelClass(id="t11", new_sample_rate=0.1, count=1)]
 
     # store our desired transaction rate
     set_transactions_resampling_rates(
@@ -68,7 +68,7 @@ def test_resampling_rates_missing():
     other_proj_id = 20
 
     other_rate = 0.1
-    other_rates = [DSElement(id="t11", count=1, new_sample_rate=0.1)]
+    other_rates = [ModelClass(id="t11", count=1, new_sample_rate=0.1)]
 
     # store some garbage to check we don't accidentally return other keys
     set_transactions_resampling_rates(


Code review comments:
diff_hunk: @@ -0,0 +1,136 @@
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+from sentry.dynamic_sampling.models.base import Model, ModelInput, ModelType
+from sentry.dynamic_sampling.models.common import ModelClass, sum_classes_counts
+from sentry.dynamic_sampling.models.full_rebalancing import FullRebalancingInput
+
+
+@dataclass
+class TransactionsRebalancingInput(ModelInput):
+    classes: List[ModelClass]
+    sample_rate: float
+    total_num_classes: Optional[int]
+    total: Optional[float]
+    intensity: float
+
+    def validate(self) -> bool:
+        return (
+            0.0 <= self.sample_rate <= 1.0
+            and 0.0 <= self.intensity <= 1.0
+            and len(self.classes) > 0
+        )
+
+
+class TransactionsRebalancingModel(
+    Model[TransactionsRebalancingInput, Tuple[List[ModelClass], float]]
+):
+    def _run(self, model_input: TransactionsRebalancingInput) -> Tuple[List[ModelClass], float]:
+        """
+        Adjusts sampling rates to bring the number of samples kept in each class as close to
+        the same value as possible while maintaining the overall sampling rate.
+
+        The algorithm adjusts the explicitly given classes individually to bring them to
+        the ideal sample rate and then adjusts the global sample rate for all the remaining classes.
+
+        :param model_input.classes: a list of class id, num_samples in class
+        :param model_input.sample_rate: global rate of sampling desired
+        :param model_input.total_num_classes: total number of classes (including the explicitly specified in classes)
+        :param model_input.intensity: the adjustment strength 0: no adjustment, 1: try to bring everything to mean
+        :param model_input.total: total number of samples in all classes (including the explicitly specified classes)
+
+        :return: a list of DSElement with calculated sample_rates and  a rate for all other (unspecified) classes.
path: src/sentry/dynamic_sampling/models/transactions_rebalancing.py
text: @user1:
nit: update with the new name for `DSElement`

diff_hunk: @@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional, Union
+
+import sentry_sdk
+
+from sentry.dynamic_sampling.models.base import Model, ModelInput
+from sentry.dynamic_sampling.rules.utils import ProjectId, TransactionName
+
+
+@dataclass
+class ModelClass:
path: src/sentry/dynamic_sampling/models/common.py
text: @user1:
I would try to avoid using the word "class" in a class name as much as possible. Judging by its usage, `RebalancedItem` (maybe with a comment explaining that item can be a project or a transaction) would be more appropriate.