Hierarchical accountant

pipeline_dp/aggregate_params.py

@@ -59,9 +59,9 @@ class AggregateParams:
     metrics: Iterable[Metrics]
     max_partitions_contributed: int
     max_contributions_per_partition: int
+    budget_weight: float = 1
     low: float = None
     high: float = None
-    budget_weight: float = 1
     public_partitions: Any = None
 
 
@@ -90,7 +90,6 @@ class SumParams:
     max_contributions_per_partition: int
     low: float
     high: float
-    budget_weight: float
     public_partitions: Union[list, 'PCollection', 'RDD']
     partition_extractor: Callable
     value_extractor: Callable

pipeline_dp/budget_accounting.py

@@ -90,6 +90,10 @@ class MechanismSpecInternal:
 class BudgetAccountant(abc.ABC):
     """Base class for budget accountants."""
 
+    def __init__(self):
+        self._scopes_stack = []
+        self._mechanisms = []
+
     @abc.abstractmethod
     def request_budget(
             self,
@@ -104,6 +108,69 @@ def request_budget(
     def compute_budgets(self):
         pass
 
+    def scope(self, weight: float):
+        """Defines a scope for DP operations that should consume no more than "weight" proportion of the budget
+        of the parent scope.
+
+        The accountant will automatically scale the budgets of all sub-operations accordingly.
+
+        Example usage:
+          with accountant.scope(weight = 0.5):
+             ... some code that consumes DP budget ...
+
+        Args:
+            weight: budget weight of all operations made within this scope as compared to.
+
+        Returns:
+            the scope that should be used in a "with" block enclosing the operations consuming the budget.
+        """
+        return BudgetAccountantScope(self, weight)
+
+    def _register_mechanism(self, mechanism: MechanismSpecInternal):
+        """Registers this mechanism for the future normalisation."""
+
+        # Register in the global list of mechanisms
+        self._mechanisms.append(mechanism)
+
+        # Register in all of the current scopes
+        for scope in self._scopes_stack:
+            scope.mechanisms.append(mechanism)
+
+        return mechanism
+
+    def _enter_scope(self, scope):
+        self._scopes_stack.append(scope)
+
+    def _exit_scope(self):
+        self._scopes_stack.pop()
+
+
+@dataclass
+class BudgetAccountantScope:
+
+    def __init__(self, accountant, weight):
+        self.weight = weight
+        self.accountant = accountant
+        self.mechanisms = []
+
+    def __enter__(self):
+        self.accountant._enter_scope(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.accountant._exit_scope()
+        self._normalise_mechanism_weights()
+
+    def _normalise_mechanism_weights(self):
+        """Normalise all mechanism weights so that they sum up to the weight of the current scope."""
+
+        if not self.mechanisms:
+            return
+
+        total_weight = sum([m.weight for m in self.mechanisms])
+        normalisation_factor = self.weight / total_weight
+        for mechanism in self.mechanisms:
+            mechanism.weight *= normalisation_factor
+
 
 class NaiveBudgetAccountant(BudgetAccountant):
     """Manages the privacy budget."""
@@ -118,12 +185,12 @@ def __init__(self, total_epsilon: float, total_delta: float):
         Raises:
             A ValueError if either argument is out of range.
         """
+        super().__init__()
 
         _validate_epsilon_delta(total_epsilon, total_delta)
 
         self._total_epsilon = total_epsilon
         self._total_delta = total_delta
-        self._mechanisms = []
 
     def request_budget(
             self,
@@ -162,7 +229,8 @@ def request_budget(
             mechanism_spec=mechanism_spec,
             sensitivity=sensitivity,
             weight=weight)
-        self._mechanisms.append(mechanism_spec_internal)
+
+        self._register_mechanism(mechanism_spec_internal)
         return mechanism_spec
 
     def compute_budgets(self):
@@ -171,6 +239,10 @@ def compute_budgets(self):
             logging.warning("No budgets were requested.")
             return
 
+        if self._scopes_stack:
+            raise Exception(
+                "Cannot call compute_budgets from within a budget scope.")
+
         total_weight_eps = total_weight_delta = 0
         for mechanism in self._mechanisms:
             total_weight_eps += mechanism.weight * mechanism.mechanism_spec.count
@@ -213,11 +285,12 @@ def __init__(self,
             ValueError: Arguments are missing or out of range.
         """
 
+        super().__init__()
+
         _validate_epsilon_delta(total_epsilon, total_delta)
 
         self._total_epsilon = total_epsilon
         self._total_delta = total_delta
-        self._mechanisms = []
         self.minimum_noise_std = None
         self._pld_discretization = pld_discretization
 
@@ -258,7 +331,7 @@ def request_budget(
             mechanism_spec=mechanism_spec,
             sensitivity=sensitivity,
             weight=weight)
-        self._mechanisms.append(mechanism_spec_internal)
+        self._register_mechanism(mechanism_spec_internal)
         return mechanism_spec
 
     def compute_budgets(self):
@@ -269,7 +342,13 @@ def compute_budgets(self):
         entire pipeline.
         """
         if not self._mechanisms:
+            logging.warning("No budgets were requested.")
             return
+
+        if self._scopes_stack:
+            raise Exception(
+                "Cannot call compute_budgets from within a budget scope.")
+
         if self._total_delta == 0:
             sum_weights = 0
             for mechanism in self._mechanisms:

pipeline_dp/dp_engine.py

@@ -52,6 +52,13 @@ def aggregate(self, col, params: AggregateParams,
         """
         if params is None:
             return None
+
+        with self._budget_accountant.scope(weight=params.budget_weight):
+            return self._aggregate(col, params, data_extractors)
+
+    def _aggregate(self, col, params: AggregateParams,
+                   data_extractors: DataExtractors):
+
         self._report_generators.append(ReportGenerator(params))
 
         accumulator_factory = AccumulatorFactory(
@@ -210,4 +217,4 @@ def _fix_budget_accounting_if_needed(self, col, accumulator_factory):
             return col
         mechanism_specs = accumulator_factory.get_mechanism_specs()
         return self._ops.map_values(
-            col, lambda acc: acc.set_mechanism_specs(mechanism_specs))
+            col, lambda acc: acc.set_mechanism_specs(mechanism_specs))

tests/budget_accounting_test.py

@@ -55,6 +55,29 @@ def test_compute_budgets(self):
         self.assertEqual(budget2.eps, 0.75)
         self.assertEqual(budget2.delta, 1e-6)
 
+    def test_budget_scopes(self):
+        budget_accountant = NaiveBudgetAccountant(total_epsilon=1,
+                                                  total_delta=1e-6)
+
+        with budget_accountant.scope(weight=0.4):
+            budget1 = budget_accountant.request_budget(
+                mechanism_type=MechanismType.LAPLACE)
+            budget2 = budget_accountant.request_budget(
+                mechanism_type=MechanismType.LAPLACE, weight=3)
+
+        with budget_accountant.scope(weight=0.6):
+            budget3 = budget_accountant.request_budget(
+                mechanism_type=MechanismType.LAPLACE)
+            budget4 = budget_accountant.request_budget(
+                mechanism_type=MechanismType.LAPLACE, weight=4)
+
+        budget_accountant.compute_budgets()
+
+        self.assertEqual(budget1.eps, 0.4 * (1 / 4))
+        self.assertEqual(budget2.eps, 0.4 * (3 / 4))
+        self.assertEqual(budget3.eps, 0.6 * (1 / 5))
+        self.assertEqual(budget4.eps, 0.6 * (4 / 5))
+
     def test_count(self):
         budget_accountant = NaiveBudgetAccountant(total_epsilon=1,
                                                   total_delta=1e-6)