parallel_analysis_node.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pyre-strict

import logging
import os
import tempfile
from concurrent.futures import ProcessPoolExecutor
from typing import Optional

import numpy as np
import torch
from privacy_guard.analysis.base_analysis_input import BaseAnalysisInput
from privacy_guard.analysis.mia.analysis_node import AnalysisNode, AnalysisNodeOutput
from privacy_guard.analysis.mia.mia_results import MIAResults


logger: logging.Logger = logging.getLogger(__name__)

TimerStats = dict[str, float]


class ParallelAnalysisNode(AnalysisNode):
    """
    ParallelAnalysisNode class for PrivacyGuard, which is an extension of AnalysisNode to parallelize the bootstrapping process for computing epsilon CI.

    args:
        analysis_input: AnalysisInput object containing the training (members) and testing (non-members) dataframes
        delta: delta parameter in (epsilon, delta)-differential privacy, close to 0
        n_users_for_eval: number of users to use for computing epsilon with Clopper-Pearson method
        eps_computation_tasks_num: number of tasks for parallel computation
        num_bootstrap_resampling_times: number of times to resample the training and testing data for computing bootstrap confidence intervals
        use_upper_bound: boolean for whether to compute epsilon at the upper-bound of CI
        use_fnr_tnr: boolean for whether to use FNR and TNR in addition to FPR and TPR error thresholds in eps_max_array computation.
        with_timer: boolean for whether to show timer for analysis node
        tpr_target: Optional target TPR for computing epsilon. If None (default), uses legacy 1% intervals.
        tpr_threshold_width: Width between TPR thresholds for fine-grained mode. Default 0.0025.
    """

    def __init__(
        self,
        analysis_input: BaseAnalysisInput,
        delta: float,
        n_users_for_eval: int,
        eps_computation_tasks_num: int,
        use_upper_bound: bool = True,
        num_bootstrap_resampling_times: int = 1000,
        use_fnr_tnr: bool = False,
        with_timer: bool = False,
        tpr_target: Optional[float] = None,
        tpr_threshold_width: float = 0.0025,
    ) -> None:
        super().__init__(
            analysis_input=analysis_input,
            delta=delta,
            n_users_for_eval=n_users_for_eval,
            use_upper_bound=use_upper_bound,
            num_bootstrap_resampling_times=num_bootstrap_resampling_times,
            use_fnr_tnr=use_fnr_tnr,
            with_timer=with_timer,
            tpr_target=tpr_target,
            tpr_threshold_width=tpr_threshold_width,
        )
        self._eps_computation_tasks_num = eps_computation_tasks_num

    def _compute_metrics_array(
        self, args: tuple[str, str, int]
    ) -> list[tuple[float, float, list[float], list[float], list[float]]]:
        """
        Make a tuple with two lists:
        -- A list of tuples metrics at error thresholds, each of which contains the
        accuracy, AUC, and epsilon values for a given number of samples
        The tuples are randomly generated from permutations of subsets of loss_train
        and loss_test.
        -- A list of epsilons at error thresholds

        Args:
            A tuple of train/test filenames and chunk size (number of bootstrap resampling times)
        Returns:
            A List with elements (accuracy, auc_value, eps_fpr_array, eps_tpr_array, eps_max_array)
        """

        train_filename, test_filename, chunk_size = args
        loss_train: torch.Tensor = torch.load(train_filename, weights_only=True)
        loss_test: torch.Tensor = torch.load(test_filename, weights_only=True)
        train_size, test_size = loss_train.shape[0], loss_test.shape[0]
        bootstrap_sample_size = min(train_size, test_size)
        metrics_results = []

        try:
            for _ in range(chunk_size):
                mia_results = MIAResults(
                    loss_train[
                        self._compute_bootstrap_sample_indexes(
                            train_size, bootstrap_sample_size
                        )
                    ],
                    loss_test[
                        self._compute_bootstrap_sample_indexes(
                            test_size, bootstrap_sample_size
                        )
                    ],
                )

                metrics_result = mia_results.compute_metrics_at_error_threshold(
                    self._delta,
                    error_threshold=self._error_thresholds,
                    use_fnr_tnr=self._use_fnr_tnr,
                    verbose=False,
                )
                metrics_results.append(metrics_result)
        except Exception as e:
            logger.info(
                f"An exception occurred when computing acc/auc/epsilon metrics: {e}"
            )

        return metrics_results

    def _parallel_compute_chunk_sizes(self, task_num: int) -> list[int]:
        """
        Compute chunk sizes for parallel computation given a task number

        Args:
            task_num (int): number of tasks for parallel computation
        Returns:
            A list of chunk sizes
        """

        base_chunk_size = self._num_bootstrap_resampling_times // task_num
        chunk_sizes = [base_chunk_size] * task_num
        if (remainder := self._num_bootstrap_resampling_times % task_num) > 0:
            for index in range(remainder):
                chunk_sizes[index] += 1

        return chunk_sizes

    def run_analysis(self) -> AnalysisNodeOutput:
        """
        Computes analysis outputs based on the input dataframes.
        Overrides "BaseAnalysisNode::run_analysis"

        First, makes loss_train and loss_test and computes one off metrics like epsilon confidence intervals.

        Then, uses "make_metrics_array" to build lists of
        metrics, each computed from random subsets of
        loss_train and loss_test. The length of these lists
        is determined by self._num_bootstrap_resampling_times

        These metrics are combined into the output of this analysis,
        and returned from the call.

        Returns:
            AnalysisNodeOutput dataclass with fields:
            "eps": epsilon at TPR=1% UB threshold if use_upper_bound is True, else epsilon at TPR=1% LB threshold
            "eps_fpr_max_lb", "eps_fpr_lb", "eps_fpr_ub": epsilon at various false positive rate:
            "eps_tpr_lb", "eps_tpr_ub": epsilon at various true positive rates
            "eps_max_lb", "eps_max_ub": max of epsilon at various true positive rates and false positive rates
            "eps_cp": epsilon calculated with Clopper-Pearson confidence interval
            "accuracy", "accuracy_ci": accuracy values
            "auc", "auc_ci": area under ROC curve values
            "data_size": dictionary with keys "train_size", "test_size", "bootstrap_size"
        """
        score_train = self.analysis_input.df_train_user["score"]
        score_test = self.analysis_input.df_test_user["score"]
        loss_train = torch.from_numpy(score_train.to_numpy())
        loss_test = torch.from_numpy(score_test.to_numpy())
        train_size, test_size = loss_train.shape[0], loss_test.shape[0]
        bootstrap_sample_size = min(train_size, test_size)
        logger.info(f"Train/Test unique users: {train_size}/{test_size}")

        metrics_array = []

        with self.timer("parallel_bootstrap"):
            with tempfile.TemporaryDirectory() as temp_dir:
                train_filename = os.path.join(temp_dir, "train_scores.pt")
                test_filename = os.path.join(temp_dir, "test_scores.pt")
                torch.save(loss_train, train_filename)
                torch.save(loss_test, test_filename)
                chunk_sizes = self._parallel_compute_chunk_sizes(
                    self._eps_computation_tasks_num
                )

                with ProcessPoolExecutor(
                    max_workers=self._eps_computation_tasks_num
                ) as pool:
                    results = pool.map(
                        self._compute_metrics_array,
                        [
                            (
                                train_filename,
                                test_filename,
                                chunk_size,
                            )
                            for chunk_size in chunk_sizes
                        ],
                    )

                for metrics_result in results:
                    metrics_array.extend(metrics_result)

        accuracy = np.array([run[0] for run in metrics_array])
        auc = np.array([run[1] for run in metrics_array])
        eps_fpr = np.array([run[2] for run in metrics_array])
        eps_tpr = np.array([run[3] for run in metrics_array])
        eps_max = np.array([run[4] for run in metrics_array])

        logger.info(
            f"epsilon at TPR thresholds: eps_tpr_array shape {eps_tpr.shape} - has NaNs: {np.isnan(eps_tpr).any()}"
        )

        # get CI bounds with 95% confidence
        accuracy_lb, accuracy_ub = self._compute_ci(accuracy)
        auc_lb, auc_ub = self._compute_ci(auc)
        eps_fpr_lb, eps_fpr_ub = self._compute_ci(eps_fpr)
        eps_tpr_lb, eps_tpr_ub = self._compute_ci(eps_tpr)
        eps_max_lb, eps_max_ub = self._compute_ci(eps_max)

        eps_tpr_boundary = eps_tpr_ub if self._use_upper_bound else eps_tpr_lb

        tpr_idx = self._get_tpr_index()
        outputs = AnalysisNodeOutput(
            eps=eps_tpr_boundary[tpr_idx],  # epsilon at specified TPR threshold
            eps_lb=eps_tpr_lb[tpr_idx],
            eps_fpr_max_ub=np.nanmax(eps_fpr_ub),
            eps_fpr_lb=list(eps_fpr_lb),
            eps_fpr_ub=list(eps_fpr_ub),
            eps_tpr_lb=list(eps_tpr_lb),
            eps_tpr_ub=list(eps_tpr_ub),
            eps_max_lb=list(eps_max_lb),
            eps_max_ub=list(eps_max_ub),
            eps_cp=float(
                "nan"
            ),  # TODO: compute eps_cp properly, currently not computed in ParallelAnalysisNode
            accuracy=accuracy.mean(),
            accuracy_ci=[accuracy_lb[0], accuracy_ub[0]],
            auc=auc.mean(),
            auc_ci=[auc_lb[0], auc_ub[0]],
            data_size={
                "train_size": train_size,
                "test_size": test_size,
                "bootstrap_size": bootstrap_sample_size,
            },
            tpr_target=self._tpr_target,
            tpr_idx=tpr_idx,
        )

        if self._with_timer:
            logger.info(f"Timer stats: {self.get_timer_stats()}")

        return outputs