Return budget allocation details from BudgetAccountant

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget/BUILD.bazel

@@ -25,6 +25,11 @@ package(
     ],
 )
 
+kt_jvm_library(
+    name = "budget_allocation_details",
+    srcs = ["BudgetAllocationDetails.kt"],
+)
+
 kt_jvm_library(
     name = "budget_spec",
     srcs = ["BudgetSpec.kt"],
@@ -43,6 +48,8 @@ kt_jvm_library(
     srcs = ["BudgetAccountant.kt"],
     deps = [
         ":allocated_budget",
+        ":budget_allocation_details",
         ":budget_spec",
+        "@maven//:com_google_errorprone_error_prone_annotations",
     ],
 )

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget/BudgetAccountant.kt

@@ -16,6 +16,7 @@
 
 package com.google.privacy.differentialprivacy.pipelinedp4j.core.budget
 
+import com.google.errorprone.annotations.CanIgnoreReturnValue
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetAccountingStrategy.NAIVE
 import java.lang.IllegalArgumentException
 import java.lang.IllegalStateException
@@ -40,12 +41,15 @@ interface BudgetAccountant {
   fun requestBudget(budgetRequest: BudgetRequest): AllocatedBudget
 
   /**
-   * Allocates budgets to all previously recorded [BudgetRequest]s. This method should only be
-   * called once.
+   * Allocates budgets to all previously recorded [BudgetRequest]s.
    *
+   * This method should only be called once.
+   *
+   * @return the allocation details for each requested budget (e.g. very useful to understand how
+   *   much budget was allocated to relative budget requests).
    * @throws IllegalStateException if budgets have already been allocated.
    */
-  fun allocateBudgets()
+  @CanIgnoreReturnValue fun allocateBudgets(): List<BudgetAllocationDetails>
 }
 
 /**
@@ -100,7 +104,8 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     return allocatedBudget
   }
 
-  override fun allocateBudgets() {
+  @CanIgnoreReturnValue
+  override fun allocateBudgets(): List<BudgetAllocationDetails> {
     if (budgetsAllocated) {
       throw IllegalStateException("Budgets have already been allocated.")
     }
@@ -119,8 +124,11 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     checkEnoughAbsoluteBudget(totalRequestedEpsilon, totalRequestedDelta)
     checkEnoughRelativeBudget(remainingEpsilon, remainingDelta)
 
-    allocateAbsoluteBudgets()
-    allocateRelativeBudgets(remainingEpsilon, remainingDelta)
+    initializeAbsoluteBudgets()
+    initializeRelativeBudgets(remainingEpsilon, remainingDelta)
+
+    return absoluteBudgets.map { it.toBudgetAllocationDetails() } +
+      relativeBudgets.map { it.toBudgetAllocationDetails() }
   }
 
   private fun checkEnoughAbsoluteBudget(requestedEpsilon: Double, requestedDelta: Double) {
@@ -147,7 +155,7 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     return Math.abs(diff) > remaining / FLOATING_POINT_ARITHMETICS_TOLERANCE
   }
 
-  private fun allocateAbsoluteBudgets() {
+  private fun initializeAbsoluteBudgets() {
     for (requestedAndAllocated in absoluteBudgets) {
       val budgetSpec = requestedAndAllocated.requested.budgetSpec as AbsoluteBudgetPerOpSpec
       requestedAndAllocated.allocated.initialize(
@@ -157,7 +165,7 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     }
   }
 
-  private fun allocateRelativeBudgets(remainingEpsilon: Double, remainingDelta: Double) {
+  private fun initializeRelativeBudgets(remainingEpsilon: Double, remainingDelta: Double) {
     var totalEpsilonWeight = 0.0
     var totalDeltaWeight = 0.0
     for (requestedAndAllocated in relativeBudgets) {
@@ -200,11 +208,13 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     }
   }
 
-  private fun relativeEpsilonRequested(): Boolean =
-    relativeBudgets.any { it.requested.mechanism.usesEpsilon }
+  private fun relativeEpsilonRequested(): Boolean = relativeBudgets.any {
+    it.requested.mechanism.usesEpsilon
+  }
 
-  private fun relativeDeltaRequested(): Boolean =
-    relativeBudgets.any { it.requested.mechanism.usesDelta }
+  private fun relativeDeltaRequested(): Boolean = relativeBudgets.any {
+    it.requested.mechanism.usesDelta
+  }
 }
 
 /**
@@ -243,7 +253,23 @@ enum class AccountedMechanism(val usesEpsilon: Boolean, val usesDelta: Boolean)
 internal data class RequestedAndAllocatedBudget(
   val requested: BudgetRequest,
   val allocated: AllocatedBudget,
-)
+) {
+  /** Converts a [RequestedAndAllocatedBudget] to a [BudgetAllocationDetails]. */
+  fun toBudgetAllocationDetails(): BudgetAllocationDetails {
+    val epsilon = allocated.epsilon()
+    val delta = allocated.delta()
+    return when (requested.mechanism) {
+      AccountedMechanism.GAUSSIAN_NOISE ->
+        BudgetAllocationDetails.GaussianAggregationAllocation(epsilon, delta)
+      AccountedMechanism.LAPLACE_NOISE ->
+        BudgetAllocationDetails.LaplaceAggregationAllocation(epsilon)
+      AccountedMechanism.PREAGGREGATED_PARTITION_SELECTION ->
+        BudgetAllocationDetails.PreaggregatedPartitionSelectionAllocation(epsilon, delta)
+      AccountedMechanism.POSTAGGREGATED_PARTITION_SELECTION ->
+        BudgetAllocationDetails.PostaggregatedPartitionSelectionAllocation(delta)
+    }
+  }
+}
 
 enum class BudgetAccountingStrategy {
   NAIVE

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget/BudgetAllocationDetails.kt

@@ -0,0 +1,46 @@
+package com.google.privacy.differentialprivacy.pipelinedp4j.core.budget
+
+/**
+ * Sealed class representing the details of a budget allocation for different differential privacy
+ * mechanisms.
+ *
+ * Each subclass contains only the parameters relevant to that mechanism.
+ *
+ * This class is used to pass budget allocation details from [DpEngine] to API implementations
+ * (e.g., BeamApi), but it is not intended for direct use by end-users of the library.
+ *
+ * Extend this class if you need to propagate more details about the budget allocation from DPEngine
+ * to the backend-specific API implementations in the API package (e.g. BeamApi, etc.).
+ */
+sealed class BudgetAllocationDetails {
+  /**
+   * Budget allocation details for Gaussian mechanism used for aggregation.
+   *
+   * Uses both epsilon and delta.
+   */
+  data class GaussianAggregationAllocation(val epsilon: Double, val delta: Double) :
+    BudgetAllocationDetails()
+
+  /**
+   * Budget allocation details for Laplace mechanism used for aggregation.
+   *
+   * Uses only epsilon.
+   */
+  data class LaplaceAggregationAllocation(val epsilon: Double) : BudgetAllocationDetails()
+
+  /**
+   * Budget allocation details for pre-aggregated partition selection.
+   *
+   * Uses both epsilon and delta.
+   */
+  data class PreaggregatedPartitionSelectionAllocation(val epsilon: Double, val delta: Double) :
+    BudgetAllocationDetails()
+
+  /**
+   * Budget allocation details for post-aggregated partition selection.
+   *
+   * Only uses delta.
+   */
+  data class PostaggregatedPartitionSelectionAllocation(val delta: Double) :
+    BudgetAllocationDetails()
+}

pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget/BUILD.bazel

@@ -27,6 +27,7 @@ kt_jvm_test(
     test_class = "com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetTests",
     deps = [
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_accountant",
+        "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_allocation_details",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_spec",
         "@maven//:com_google_testparameterinjector_test_parameter_injector",
         "@maven//:com_google_truth_truth",

pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget/NaiveBudgetAccountantTest.kt

@@ -60,10 +60,15 @@ class NaiveBudgetAccountantTest {
         BudgetRequest(AbsoluteBudgetPerOpSpec(epsilon = 1.0, delta = 0.1), GAUSSIAN_NOISE)
       )
 
-    accountant.allocateBudgets()
+    val allocationDetails = accountant.allocateBudgets()
 
     assertThat(allocatedBudget.epsilon()).isEqualTo(1.0)
     assertThat(allocatedBudget.delta()).isEqualTo(0.1)
+
+    assertThat(allocationDetails)
+      .containsExactly(
+        BudgetAllocationDetails.GaussianAggregationAllocation(epsilon = 1.0, delta = 0.1)
+      )
   }
 
   @Test
@@ -110,14 +115,26 @@ class NaiveBudgetAccountantTest {
     val allocatedBudgetWeightThree =
       accountant.requestBudget(BudgetRequest(RelativeBudgetPerOpSpec(3.0), GAUSSIAN_NOISE))
 
-    accountant.allocateBudgets()
+    val allocationDetails = accountant.allocateBudgets()
 
     assertThat(allocatedBudgetWeightOne.epsilon()).isEqualTo(10.0)
     assertThat(allocatedBudgetWeightOne.delta()).isWithin(1e-13).of(0.1)
     assertThat(allocatedBudgetWeightTwo.epsilon()).isEqualTo(20.0)
     assertThat(allocatedBudgetWeightTwo.delta()).isWithin(1e-13).of(0.2)
     assertThat(allocatedBudgetWeightThree.epsilon()).isEqualTo(30.0)
     assertThat(allocatedBudgetWeightThree.delta()).isWithin(1e-13).of(0.3)
+
+    assertThat(allocationDetails).hasSize(3)
+    val details1 = allocationDetails[0] as BudgetAllocationDetails.GaussianAggregationAllocation
+    val details2 = allocationDetails[1] as BudgetAllocationDetails.GaussianAggregationAllocation
+    val details3 = allocationDetails[2] as BudgetAllocationDetails.GaussianAggregationAllocation
+
+    assertThat(details1.epsilon).isEqualTo(10.0)
+    assertThat(details1.delta).isWithin(1e-13).of(0.1)
+    assertThat(details2.epsilon).isEqualTo(20.0)
+    assertThat(details2.delta).isWithin(1e-13).of(0.2)
+    assertThat(details3.epsilon).isEqualTo(30.0)
+    assertThat(details3.delta).isWithin(1e-13).of(0.3)
   }
 
   @Test
@@ -142,6 +159,7 @@ class NaiveBudgetAccountantTest {
     val accountant = NaiveBudgetAccountant(TotalBudget(epsilon = 0.5, delta = 0.1))
     val allocatedBudget =
       accountant.requestBudget(BudgetRequest(RelativeBudgetPerOpSpec(1.0), GAUSSIAN_NOISE))
+
     accountant.allocateBudgets()
 
     assertThat(allocatedBudget.epsilon()).isEqualTo(0.5)
@@ -206,14 +224,26 @@ class NaiveBudgetAccountantTest {
     val relativeAllocatedBudgetTwiceMore =
       accountant.requestBudget(BudgetRequest(RelativeBudgetPerOpSpec(2.0), GAUSSIAN_NOISE))
 
-    accountant.allocateBudgets()
+    val allocationDetails = accountant.allocateBudgets()
 
     assertThat(absoluteAllocatedBudget.epsilon()).isEqualTo(30.0)
     assertThat(absoluteAllocatedBudget.delta()).isEqualTo(0.3)
     assertThat(relativeAllocatedBudget.epsilon()).isEqualTo(10.0)
     assertThat(relativeAllocatedBudget.delta()).isWithin(1e-13).of(0.1)
     assertThat(relativeAllocatedBudgetTwiceMore.epsilon()).isEqualTo(20.0)
     assertThat(relativeAllocatedBudgetTwiceMore.delta()).isWithin(1e-13).of(0.2)
+
+    assertThat(allocationDetails).hasSize(3)
+    val details1 = allocationDetails[0] as BudgetAllocationDetails.GaussianAggregationAllocation
+    val details2 = allocationDetails[1] as BudgetAllocationDetails.GaussianAggregationAllocation
+    val details3 = allocationDetails[2] as BudgetAllocationDetails.GaussianAggregationAllocation
+
+    assertThat(details1.epsilon).isEqualTo(30.0)
+    assertThat(details1.delta).isEqualTo(0.3)
+    assertThat(details2.epsilon).isEqualTo(10.0)
+    assertThat(details2.delta).isWithin(1e-13).of(0.1)
+    assertThat(details3.epsilon).isEqualTo(20.0)
+    assertThat(details3.delta).isWithin(1e-13).of(0.2)
   }
 
   @Test
@@ -248,13 +278,25 @@ class NaiveBudgetAccountantTest {
         BudgetRequest(RelativeBudgetPerOpSpec(2.0), POSTAGGREGATED_PARTITION_SELECTION)
       )
 
-    accountant.allocateBudgets()
+    val allocationDetails = accountant.allocateBudgets()
 
     assertThat(absoluteAllocatedBudget.epsilon()).isEqualTo(3.0)
     assertThat(absoluteAllocatedBudget.delta()).isEqualTo(0.1)
     assertThat(relativeAllocatedBudget.epsilon()).isEqualTo(7.0)
     assertThat(relativeAllocatedBudget.delta()).isWithin(1e-13).of(0.1)
     assertThat(relativeAllocatedBudgetTwiceMore.epsilon()).isEqualTo(0.0)
     assertThat(relativeAllocatedBudgetTwiceMore.delta()).isWithin(1e-13).of(0.2)
+
+    assertThat(allocationDetails).hasSize(3)
+    val details1 = allocationDetails[0] as BudgetAllocationDetails.GaussianAggregationAllocation
+    val details2 = allocationDetails[1] as BudgetAllocationDetails.GaussianAggregationAllocation
+    val details3 =
+      allocationDetails[2] as BudgetAllocationDetails.PostaggregatedPartitionSelectionAllocation
+
+    assertThat(details1.epsilon).isEqualTo(3.0)
+    assertThat(details1.delta).isEqualTo(0.1)
+    assertThat(details2.epsilon).isEqualTo(7.0)
+    assertThat(details2.delta).isWithin(1e-13).of(0.1)
+    assertThat(details3.delta).isWithin(1e-13).of(0.2)
   }
 }