def test_update(self) -> None: instance_id = self._get_random_id() test_update_private_computation_instance = PrivateComputationInstance( instance_id=instance_id, role=PrivateComputationRole.PUBLISHER, instances=[self.test_mpc_instance], status=PrivateComputationInstanceStatus.CREATED, status_update_ts=1600000000, num_files_per_mpc_container=40, game_type=PrivateComputationGameType.LIFT, input_path="in", output_dir="out", num_pid_containers=4, num_mpc_containers=4, concurrency=1, ) # Create a new MPC instance to be added to instances self.repo.create(test_update_private_computation_instance) test_mpc_instance_new = PCSMPCInstance.create_instance( instance_id=instance_id, game_name="aggregation", mpc_party=MPCParty.SERVER, num_workers=1, ) instances_new = [self.test_mpc_instance, test_mpc_instance_new] # Update instances test_update_private_computation_instance.instances = instances_new self.repo.update(test_update_private_computation_instance) # Assert instances is updated self.assertEqual(self.repo.read(instance_id).instances, instances_new) self.repo.delete(instance_id)
def test_cancel_current_stage(self) -> None: test_mpc_id = self.test_private_computation_id + "_compute_metrics" test_game_name = GameNames.LIFT.value test_mpc_party = MPCParty.CLIENT # prepare the pl instance that will be read in to memory from the repository # at the beginning of the cancel_current_stage function mpc_instance_started = PCSMPCInstance.create_instance( instance_id=test_mpc_id, game_name=test_game_name, mpc_party=test_mpc_party, num_workers=self.test_num_containers, status=MPCInstanceStatus.STARTED, ) private_computation_instance = self.create_sample_instance( status=PrivateComputationInstanceStatus.COMPUTATION_STARTED, role=PrivateComputationRole.PARTNER, instances=[mpc_instance_started], ) self.private_computation_service.instance_repository.read = MagicMock( return_value=private_computation_instance) # prepare the mpc instance that's returned from mpc_service.stop_instance() mpc_instance_canceled = PCSMPCInstance.create_instance( instance_id=test_mpc_id, game_name=test_game_name, mpc_party=test_mpc_party, num_workers=self.test_num_containers, status=MPCInstanceStatus.CANCELED, ) self.private_computation_service.mpc_svc.stop_instance = MagicMock( return_value=mpc_instance_canceled) self.private_computation_service.mpc_svc.instance_repository.read = MagicMock( return_value=mpc_instance_canceled) # call cancel, expect no exception private_computation_instance = ( self.private_computation_service.cancel_current_stage( instance_id=self.test_private_computation_id, )) # assert the pl instance returned has the correct status self.assertEqual( PrivateComputationInstanceStatus.COMPUTATION_FAILED, private_computation_instance.status, )
def setUp(self) -> None: instance_id = self._get_random_id() self.repo = LocalPrivateComputationInstanceRepository("./") self.test_mpc_instance = PCSMPCInstance.create_instance( instance_id=instance_id, game_name="conversion_lift", mpc_party=MPCParty.SERVER, num_workers=2, )
def setUp(self): self.mpc_instance = PCSMPCInstance.create_instance( instance_id=TEST_INSTANCE_ID, game_name=TEST_GAME_NAME, mpc_party=TEST_MPC_PARTY, num_workers=TEST_NUM_WORKERS, server_ips=TEST_SERVER_IPS, status=MPCInstanceStatus.CREATED, game_args=TEST_GAME_ARGS, ) self.local_instance_repo = LocalInstanceRepository(TEST_BASE_DIR)
def test_get_status_from_stage(self) -> None: # Test get status from an MPC stage mpc_instance = PCSMPCInstance.create_instance( instance_id="test_mpc_id", game_name=GameNames.SHARD_AGGREGATOR.value, mpc_party=MPCParty.SERVER, num_workers=2, status=MPCInstanceStatus.FAILED, ) pc_instance = self.create_sample_instance( PrivateComputationInstanceStatus.AGGREGATION_STARTED, instances=[mpc_instance], ) self.private_computation_service.mpc_svc.update_instance = MagicMock( return_value=mpc_instance) self.assertEqual( PrivateComputationInstanceStatus.AGGREGATION_FAILED, self.private_computation_service._update_instance( pc_instance).status, ) # Test get status from the PID stage pid_instance = PIDInstance( instance_id="test_pid_id", protocol=DEFAULT_PID_PROTOCOL, pid_role=PIDRole.PUBLISHER, num_shards=4, input_path="input", output_path="output", stages_containers={}, stages_status={ UnionPIDStage.PUBLISHER_RUN_PID: PIDStageStatus.COMPLETED }, current_stage=UnionPIDStage.PUBLISHER_RUN_PID, status=PIDInstanceStatus.COMPLETED, ) pc_instance = self.create_sample_instance( PrivateComputationInstanceStatus.ID_MATCHING_STARTED, instances=[pid_instance], ) self.private_computation_service.pid_svc.update_instance = MagicMock( return_value=pid_instance) self.assertEqual( PrivateComputationInstanceStatus.ID_MATCHING_COMPLETED, self.private_computation_service._update_instance( pc_instance).status, )
async def test_attribution_stage(self) -> None: private_computation_instance = self._create_pc_instance() mpc_instance = PCSMPCInstance.create_instance( instance_id=private_computation_instance.instance_id + "_decoupled_attribution0", game_name=GameNames.DECOUPLED_ATTRIBUTION.value, mpc_party=MPCParty.CLIENT, num_workers=private_computation_instance.num_mpc_containers, ) self.mock_mpc_svc.start_instance_async = AsyncMock(return_value=mpc_instance) test_server_ips = [ f"192.0.2.{i}" for i in range(private_computation_instance.num_mpc_containers) ] await self.stage_svc.run_async(private_computation_instance, test_server_ips) self.assertEqual(mpc_instance, private_computation_instance.instances[0])
def gen_dummy_mpc_instance() -> PCSMPCInstance: """Creates a dummy mpc instance to be used in unit tests""" return PCSMPCInstance.create_instance( instance_id="mpc_instance_id", game_name="lift", mpc_party=MPCParty.SERVER, num_workers=1, server_ips=["10.0.10.242"], containers=[gen_dummy_container_instance()], status=MPCInstanceStatus.COMPLETED, game_args=[{ "input_base_path": "https://bucket.s3.us-west-2.amazonaws.com/lift/partner/partner_instance_1638998680_0_out_dir/data_processing_stage/out.csv", "output_base_path": "https://bucket.s3.us-west-2.amazonaws.com/lift/partner/partner_instance_1638998680_0_out_dir/compute_stage/out.json", "num_files": 40, "concurrency": 4, "file_start_index": 0, }], )
async def test_compute_metrics(self) -> None: private_computation_instance = self._create_pc_instance() mpc_instance = PCSMPCInstance.create_instance( instance_id=private_computation_instance.instance_id + "_pcf2_lift0", game_name=GameNames.PCF2_LIFT.value, mpc_party=MPCParty.CLIENT, num_workers=private_computation_instance.num_mpc_containers, ) self.mock_mpc_svc.start_instance_async = AsyncMock( return_value=mpc_instance) test_server_ips = [ f"192.0.2.{i}" for i in range(private_computation_instance.num_mpc_containers) ] await self.stage_svc.run_async(private_computation_instance, test_server_ips) self.assertEqual(mpc_instance, private_computation_instance.instances[0])
def test_update_instance(self) -> None: test_pid_id = self.test_private_computation_id + "_id_match" test_pid_role = PIDRole.PUBLISHER test_input_path = "pid_in" test_output_path = "pid_out" # create one PID instance to be put into PrivateComputationInstance pid_instance = PIDInstance( instance_id=test_pid_id, protocol=DEFAULT_PID_PROTOCOL, pid_role=test_pid_role, num_shards=self.test_num_containers, input_path=test_input_path, output_path=test_output_path, status=PIDInstanceStatus.STARTED, ) private_computation_instance = self.create_sample_instance( status=PrivateComputationInstanceStatus.ID_MATCHING_STARTED, instances=[pid_instance], ) updated_pid_instance = pid_instance updated_pid_instance.status = PIDInstanceStatus.COMPLETED updated_pid_instance.current_stage = UnionPIDStage.PUBLISHER_RUN_PID updated_pid_instance.stages_status = { UnionPIDStage.PUBLISHER_RUN_PID: PIDStageStatus.COMPLETED } self.private_computation_service.pid_svc.update_instance = MagicMock( return_value=updated_pid_instance) self.private_computation_service.instance_repository.read = MagicMock( return_value=private_computation_instance) # end_ts should not be calculated until the instance run is complete. self.assertEqual(0, private_computation_instance.end_ts) # call update on the PrivateComputationInstance updated_instance = self.private_computation_service.update_instance( instance_id=self.test_private_computation_id) # check update instance called on the right pid instance # pyre-fixme[16]: Callable `update_instance` has no attribute `assert_called`. self.private_computation_service.pid_svc.update_instance.assert_called( ) self.assertEqual( test_pid_id, # pyre-fixme[16]: Callable `update_instance` has no attribute `call_args`. self.private_computation_service.pid_svc.update_instance. call_args[0][0], ) # check update instance called on the right private lift instance # pyre-fixme[16]: Callable `update` has no attribute `assert_called`. self.private_computation_service.instance_repository.update.assert_called( ) self.assertEqual( private_computation_instance, # pyre-fixme[16]: Callable `update` has no attribute `call_args`. self.private_computation_service.instance_repository.update. call_args[0][0], ) # check updated_instance has new status self.assertEqual( PrivateComputationInstanceStatus.ID_MATCHING_COMPLETED, updated_instance.status, ) # create one MPC instance to be put into PrivateComputationInstance test_mpc_id = "test_mpc_id" mpc_instance = PCSMPCInstance.create_instance( instance_id=test_mpc_id, game_name=GameNames.LIFT.value, mpc_party=MPCParty.SERVER, num_workers=2, ) private_computation_instance = self.create_sample_instance( status=PrivateComputationInstanceStatus.COMPUTATION_STARTED, instances=[mpc_instance], ) updated_mpc_instance = mpc_instance updated_mpc_instance.status = MPCInstanceStatus.COMPLETED self.private_computation_service.mpc_svc.update_instance = MagicMock( return_value=updated_mpc_instance) self.private_computation_service.instance_repository.read = MagicMock( return_value=private_computation_instance) # call update on the PrivateComputationInstance updated_instance = self.private_computation_service.update_instance( instance_id=self.test_private_computation_id) # check update instance called on the right mpc instance # pyre-fixme[16]: Callable `update_instance` has no attribute `assert_called`. self.private_computation_service.mpc_svc.update_instance.assert_called( ) self.assertEqual( test_mpc_id, # pyre-fixme[16]: Callable `update_instance` has no attribute `call_args`. self.private_computation_service.mpc_svc.update_instance. call_args[0][0], ) # check update instance called on the right private lift instance self.private_computation_service.instance_repository.update.assert_called( ) self.assertEqual( private_computation_instance, self.private_computation_service.instance_repository.update. call_args[0][0], ) # check updated_instance has new status self.assertEqual( PrivateComputationInstanceStatus.COMPUTATION_COMPLETED, updated_instance.status, ) # elapsed_time should report current running time if the run is incomplete. self.assertEqual( time.time() - private_computation_instance.creation_ts + 1, private_computation_instance.elapsed_time, ) expected_end_ts = time.time() + 2 private_computation_instance.update_status( private_computation_instance.stage_flow.get_last_stage(). completed_status, logging.getLogger(), ) self.assertEqual(expected_end_ts, private_computation_instance.end_ts) expected_elapsed_time = (private_computation_instance.end_ts - private_computation_instance.creation_ts) self.assertEqual( expected_elapsed_time, private_computation_instance.elapsed_time, )