async def run_async( self, pc_instance: PrivateComputationInstance, server_ips: Optional[List[str]] = None, ) -> PrivateComputationInstance: """Runs the pcf2.0 based private aggregation stage Args: pc_instance: the private computation instance to run aggregation stage server_ips: only used by the partner role. These are the ip addresses of the publisher's containers. Returns: An updated version of pc_instance that stores an MPCInstance """ # Prepare arguments for attribution game game_args = self._get_compute_metrics_game_args(pc_instance, ) # We do this check here because depends on how game_args is generated, len(game_args) could be different, # but we will always expect server_ips == len(game_args) if server_ips and len(server_ips) != len(game_args): raise ValueError( f"Unable to rerun MPC pcf2.0 based aggregation because there is a mismatch between the number of server ips given ({len(server_ips)}) and the number of containers ({len(game_args)}) to be spawned." ) # Create and start MPC instance to run MPC compute logging.info( "Starting to run MPC instance for pcf2.0 based aggregation stage.") stage_data = PrivateComputationServiceData.PCF2_AGGREGATION_STAGE_DATA binary_name = OneDockerBinaryNames.PCF2_AGGREGATION.value game_name = checked_cast(str, stage_data.game_name) binary_config = self._onedocker_binary_config_map[binary_name] retry_counter_str = str(pc_instance.retry_counter) mpc_instance = await create_and_start_mpc_instance( mpc_svc=self._mpc_service, instance_id=pc_instance.instance_id + "_" + GameNames.PCF2_AGGREGATION.value + retry_counter_str, game_name=game_name, mpc_party=map_private_computation_role_to_mpc_party( pc_instance.role), num_containers=len(game_args), binary_version=binary_config.binary_version, server_ips=server_ips, game_args=game_args, container_timeout=self._container_timeout, repository_path=binary_config.repository_path, ) logging.info( "MPC instance started running for pcf2.0 based aggregation stage.") # Push MPC instance to PrivateComputationInstance.instances and update PL Instance status pc_instance.instances.append( PCSMPCInstance.from_mpc_instance(mpc_instance)) return pc_instance
def _get_attribution_game_args( self, private_computation_instance: PrivateComputationInstance, common_compute_game_args: Dict[str, Any], ) -> List[Dict[str, Any]]: """Gets attribution specific game args to be passed to game binaries by onedocker When onedocker spins up containers to run games, it unpacks a dictionary containing the arguments required by the game binary being ran. This function prepares arguments specific to attribution games. Args: pc_instance: the private computation instance to generate game args for Returns: MPC game args to be used by onedocker """ game_args = [] aggregation_type = checked_cast( AggregationType, private_computation_instance.aggregation_type) attribution_rule = checked_cast( AttributionRule, private_computation_instance.attribution_rule) game_args = [{ **common_compute_game_args, **{ "aggregators": aggregation_type.value, "attribution_rules": attribution_rule.value, "file_start_index": i * private_computation_instance.num_files_per_mpc_container, "use_xor_encryption": True, "run_name": private_computation_instance.instance_id if self._log_cost_to_s3 else "", "max_num_touchpoints": private_computation_instance.padding_size, "max_num_conversions": private_computation_instance.padding_size, }, } for i in range(private_computation_instance.num_mpc_containers)] return game_args
def _get_compute_metrics_game_args( self, private_computation_instance: PrivateComputationInstance, ) -> List[Dict[str, Any]]: """Gets the game args passed to game binaries by onedocker When onedocker spins up containers to run games, it unpacks a dictionary containing the arguments required by the game binary being ran. This function prepares that dictionary. Args: pc_instance: the private computation instance to generate game args for Returns: MPC game args to be used by onedocker """ attribution_rule = checked_cast( AttributionRule, private_computation_instance.attribution_rule) if self._log_cost_to_s3: run_name = (private_computation_instance.instance_id + "_" + GameNames.PCF2_ATTRIBUTION.value) if private_computation_instance.post_processing_data: private_computation_instance.post_processing_data.s3_cost_export_output_paths.add( f"att-logs/{run_name}_{private_computation_instance.role.value.title()}.json" ) else: run_name = "" common_game_args = { "input_base_path": private_computation_instance.data_processing_output_path, "output_base_path": private_computation_instance. pcf2_attribution_stage_output_base_path, "num_files": private_computation_instance.num_files_per_mpc_container, "concurrency": private_computation_instance.concurrency, "run_name": run_name, "max_num_touchpoints": private_computation_instance.padding_size, "max_num_conversions": private_computation_instance.padding_size, "log_cost": self._log_cost_to_s3, "attribution_rules": attribution_rule.value, "use_xor_encryption": True, "use_postfix": True, } game_args = [{ **common_game_args, **{ "file_start_index": i * private_computation_instance.num_files_per_mpc_container, }, } for i in range(private_computation_instance.num_mpc_containers)] return game_args
async def run( self, stage_input: PIDStageInput, container_timeout: Optional[int] = None, wait_for_containers: bool = True, ) -> PIDStageStatus: self.logger.info(f"[{self}] Called run") instance_id = stage_input.instance_id timeout = container_timeout or DEFAULT_CONTAINER_TIMEOUT_IN_SEC # Make sure status is READY before proceed status = await self._ready(stage_input) await self.update_instance_status(instance_id=instance_id, status=status) if status is not PIDStageStatus.READY: return status # Some invariant checking on the input and output paths input_paths = stage_input.input_paths output_paths = stage_input.output_paths num_shards = (stage_input.num_shards + 1 if stage_input.is_validating else stage_input.num_shards) if len(input_paths) != 1: raise ValueError(f"Expected 1 input path, not {len(input_paths)}") if len(output_paths) != 1: raise ValueError( f"Expected 1 output path, not {len(output_paths)}") await self.update_instance_status(instance_id=instance_id, status=PIDStageStatus.STARTED) if stage_input.pid_use_row_numbers: self.logger.info("use-row-numbers is enabled for Private ID") if self.stage_type is UnionPIDStage.PUBLISHER_RUN_PID: # Run publisher commands in container self.logger.info("Publisher spinning up containers") try: binary = OneDockerBinaryNames.PID_SERVER.value if self.protocol == PIDProtocol.UNION_PID_MULTIKEY: binary = OneDockerBinaryNames.PID_MULTI_KEY_SERVER.value pending_containers = self.onedocker_svc.start_containers( package_name=binary, version=self.onedocker_binary_config.binary_version, cmd_args_list=self._gen_command_args_list( input_path=input_paths[0], output_path=output_paths[0], num_shards=num_shards, use_row_numbers=stage_input.pid_use_row_numbers and (self.protocol != PIDProtocol.UNION_PID_MULTIKEY), disable_metric_logging=False, ), env_vars=self._gen_env_vars(), timeout=timeout, ) containers = await self.onedocker_svc.wait_for_pending_containers( [ container.instance_id for container in pending_containers ]) except Exception as e: status = PIDStageStatus.FAILED await self.update_instance_status(instance_id=instance_id, status=status) self.logger.exception(f"Failed to spin up containers: {e}") return status # Write containers information to PID instance repository await self.update_instance_containers(instance_id=instance_id, containers=containers) # Get ips from containers and write them to pid instance repository self.logger.info("Storing servers' IPs") ip_addresses = [ checked_cast(str, container.ip_address) for container in containers ] await self.put_server_ips(instance_id=instance_id, server_ips=ip_addresses) # Wait until the containers are finished if wait_for_containers: self.logger.info("Waiting for containers to finish") containers = await RunBinaryBaseService.wait_for_containers_async( self.onedocker_svc, containers, SLEEP_UPDATE_SECONDS) await self.update_instance_containers(instance_id=instance_id, containers=containers) status = self.get_stage_status_from_containers(containers) elif self.stage_type is UnionPIDStage.ADV_RUN_PID: server_ips = self.server_ips or [] if not server_ips: self.logger.error("Missing server_ips") status = PIDStageStatus.FAILED await self.update_instance_status(instance_id=instance_id, status=status) return status hostnames = [f"http://{ip}" for ip in server_ips] # Run partner commands in container self.logger.info("Partner spinning up containers") try: binary = OneDockerBinaryNames.PID_CLIENT.value if self.protocol == PIDProtocol.UNION_PID_MULTIKEY: binary = OneDockerBinaryNames.PID_MULTI_KEY_CLIENT.value pending_containers = self.onedocker_svc.start_containers( package_name=binary, version=self.onedocker_binary_config.binary_version, cmd_args_list=self._gen_command_args_list( input_path=input_paths[0], output_path=output_paths[0], num_shards=num_shards, server_hostnames=hostnames, use_row_numbers=stage_input.pid_use_row_numbers and (self.protocol != PIDProtocol.UNION_PID_MULTIKEY), disable_metric_logging=True, ), env_vars=self._gen_env_vars(), timeout=timeout, ) containers = await self.onedocker_svc.wait_for_pending_containers( [ container.instance_id for container in pending_containers ]) except Exception as e: status = PIDStageStatus.FAILED await self.update_instance_status(instance_id=instance_id, status=status) self.logger.exception(f"Failed to spin up containers: {e}") return status # Write containers information to PID instance repository await self.update_instance_containers(instance_id=instance_id, containers=containers) if wait_for_containers: # Wait until the containers are finished self.logger.info("Waiting for containers to finish") containers = await RunBinaryBaseService.wait_for_containers_async( self.onedocker_svc, containers, SLEEP_UPDATE_SECONDS) await self.update_instance_containers(instance_id=instance_id, containers=containers) status = self.get_stage_status_from_containers(containers) self.logger.info(f"PID Run protocol status: {status}") await self.update_instance_status(instance_id=instance_id, status=status) return status
async def start_combiner_service( private_computation_instance: PrivateComputationInstance, onedocker_svc: OneDockerService, onedocker_binary_config_map: DefaultDict[str, OneDockerBinaryConfig], combine_output_path: str, log_cost_to_s3: bool = DEFAULT_LOG_COST_TO_S3, wait_for_containers: bool = False, max_id_column_count: int = 1, ) -> List[ContainerInstance]: """Run combiner service and return those container instances Args: private_computation_instance: The PC instance to run combiner service with onedocker_svc: Spins up containers that run binaries in the cloud onedocker_binary_config_map: Stores a mapping from mpc game to OneDockerBinaryConfig (binary version and tmp directory) combine_output_path: out put path for the combine result log_cost_to_s3: if money cost of the computation will be logged to S3 wait_for_containers: block until containers to finish running, default False Returns: return: list of container instances running combiner service """ stage_data = PrivateComputationServiceData.get( private_computation_instance.game_type).combiner_stage binary_name = stage_data.binary_name binary_config = onedocker_binary_config_map[binary_name] # TODO: T106159008 Add on attribution specific args if private_computation_instance.game_type is PrivateComputationGameType.ATTRIBUTION: run_name = private_computation_instance.instance_id if log_cost_to_s3 else "" padding_size = checked_cast(int, private_computation_instance.padding_size) log_cost = log_cost_to_s3 else: run_name = None padding_size = None log_cost = None combiner_service = checked_cast( IdSpineCombinerService, stage_data.service, ) args = combiner_service.build_args( spine_path=private_computation_instance.pid_stage_output_spine_path, data_path=private_computation_instance.pid_stage_output_data_path, output_path=combine_output_path, num_shards=private_computation_instance.num_pid_containers + 1 if private_computation_instance.is_validating else private_computation_instance.num_pid_containers, tmp_directory=binary_config.tmp_directory, max_id_column_cnt=max_id_column_count, run_name=run_name, padding_size=padding_size, log_cost=log_cost, ) env_vars = {ONEDOCKER_REPOSITORY_PATH: binary_config.repository_path} return await combiner_service.start_containers( cmd_args_list=args, onedocker_svc=onedocker_svc, binary_version=binary_config.binary_version, binary_name=binary_name, timeout=None, wait_for_containers_to_finish=wait_for_containers, env_vars=env_vars, )
def server_ips(self) -> List[str]: return [ checked_cast(str, container.ip_address) for container in self.containers ]
def _get_compute_metrics_game_args( self, private_computation_instance: PrivateComputationInstance, ) -> List[Dict[str, Any]]: """Gets the game args passed to game binaries by onedocker When onedocker spins up containers to run games, it unpacks a dictionary containing the arguments required by the game binary being ran. This function prepares that dictionary. Args: pc_instance: the private computation instance to generate game args for Returns: MPC game args to be used by onedocker """ logging.info( f"PATH : {private_computation_instance.pcf2_attribution_stage_output_base_path}" ) aggregation_type = checked_cast( AggregationType, private_computation_instance.aggregation_type) # pcf2.0 based aggregation game does not need the attribution rule. Passing it here just for cost logging, # so that we get to know how much the game cost for a aggregation format and attribution rule. attribution_rule = checked_cast( AttributionRule, private_computation_instance.attribution_rule) if self._log_cost_to_s3: run_name = private_computation_instance.instance_id if private_computation_instance.post_processing_data: private_computation_instance.post_processing_data.s3_cost_export_output_paths.add( f"agg-logs/{run_name}_{private_computation_instance.role.value.title()}.json", ) else: run_name = "" common_game_args = { "input_base_path": private_computation_instance.data_processing_output_path, "output_base_path": private_computation_instance. pcf2_aggregation_stage_output_base_path, "num_files": private_computation_instance.num_files_per_mpc_container, "concurrency": private_computation_instance.concurrency, "aggregators": aggregation_type.value, "attribution_rules": attribution_rule.value, "input_base_path_secret_share": private_computation_instance. pcf2_attribution_stage_output_base_path, "use_xor_encryption": True, "use_postfix": True, "run_name": run_name, "max_num_touchpoints": private_computation_instance.padding_size, "max_num_conversions": private_computation_instance.padding_size, "log_cost": self._log_cost_to_s3, } game_args = [{ **common_game_args, **{ "file_start_index": i * private_computation_instance.num_files_per_mpc_container, }, } for i in range(private_computation_instance.num_mpc_containers)] return game_args