コード例 #1
0
    def setUp(self, onedocker_service) -> None:
        self.onedocker_service = onedocker_service
        self.test_num_containers = 2

        self.onedocker_binary_config_map = defaultdict(
            lambda: OneDockerBinaryConfig(
                tmp_directory="/test_tmp_directory/",
                binary_version="latest",
                repository_path="test_path/",
            ))
        self.stage_svc = PrepareDataStageService(
            self.onedocker_service, self.onedocker_binary_config_map)
コード例 #2
0
    def setUp(self, mock_mpc_svc: MPCService) -> None:
        self.mock_mpc_svc = mock_mpc_svc
        self.mock_mpc_svc.create_instance = MagicMock()

        onedocker_binary_config_map = defaultdict(
            lambda: OneDockerBinaryConfig(
                tmp_directory="/test_tmp_directory/",
                binary_version="latest",
                repository_path="test_path/",
            ))
        self.stage_svc = ComputeMetricsStageService(
            onedocker_binary_config_map, self.mock_mpc_svc)
コード例 #3
0
 def setUp(self, mock_onedocker_svc, mock_storage_svc) -> None:
     self.mock_onedocker_svc = mock_onedocker_svc
     self.mock_storage_svc = mock_storage_svc
     self.onedocker_binary_config = OneDockerBinaryConfig(
         tmp_directory="/tmp",
         binary_version="latest",
         repository_path="test_path/",
     )
     self.binary_name = "data_processing/pid_preparer"
     self.onedocker_binary_config_map = {
         self.binary_name: self.onedocker_binary_config
     }
     self.input_path = "in"
     self.output_path = "out"
     self.pc_instance_id = "test_instance_123"
     self.container_timeout = 43200
コード例 #4
0
        async def _run_sub_test(
            wait_for_containers: bool,
            expected_container_status: ContainerInstanceStatus,
        ) -> None:
            with patch.object(PIDStage, "update_instance_containers"):
                test_onedocker_binary_config = OneDockerBinaryConfig(
                    tmp_directory="/test_tmp_directory/",
                    binary_version="latest",
                    repository_path="test_path/",
                )
                container = ContainerInstance(
                    instance_id="123",
                    ip_address="192.0.2.0",
                    status=expected_container_status,
                )
                mock_sharder.return_value = [container]
                stage = PIDShardStage(
                    stage=UnionPIDStage.PUBLISHER_SHARD,
                    instance_repository=mock_instance_repo,
                    storage_svc=mock_storage_svc,
                    onedocker_svc=mock_onedocker_svc,
                    onedocker_binary_config=test_onedocker_binary_config,
                )

                test_input_path = "foo"
                test_output_path = "bar"
                test_num_shards = 1
                test_hmac_key = "CoXbp7BOEvAN9L1CB2DAORHHr3hB7wE7tpxMYm07tc0="

                shard_path = PIDShardStage.get_sharded_filepath(
                    test_output_path, 0)
                self.assertEqual(f"{test_output_path}_0", shard_path)

                res = await stage.shard(
                    "123",
                    test_input_path,
                    test_output_path,
                    test_num_shards,
                    test_hmac_key,
                    wait_for_containers=wait_for_containers,
                )
                self.assertEqual(
                    PIDStage.get_stage_status_from_containers([container]),
                    res,
                )

                mock_sharder.assert_called_once()
コード例 #5
0
 def setUp(self, mock_onedocker_service, mock_storage_service) -> None:
     self.mock_onedocker_svc = mock_onedocker_service
     self.mock_storage_svc = mock_storage_service
     self.test_num_containers = 1
     self.onedocker_binary_config_map = defaultdict(
         lambda: OneDockerBinaryConfig(
             tmp_directory="/test_tmp_directory/",
             binary_version="latest",
             repository_path="test_path/",
         )
     )
     self.server_ips = [f"192.0.2.{i}" for i in range(self.test_num_containers)]
     self.input_path = "in"
     self.output_path = "out"
     self.pc_instance_id = "test_instance_123"
     self.port = 15200
     self.use_row_numbers = True
コード例 #6
0
 def setUp(self, mock_onedocker_svc, mock_storage_svc) -> None:
     self.mock_onedocker_svc = mock_onedocker_svc
     self.mock_storage_svc = mock_storage_svc
     self.onedocker_binary_config = OneDockerBinaryConfig(
         tmp_directory="/tmp",
         binary_version="latest",
         repository_path="test_path/",
     )
     self.onedocker_binary_config = self.onedocker_binary_config
     self.binary_name = "data_processing/sharder_hashed_for_pid"
     self.onedocker_binary_config_map = {
         self.binary_name: self.onedocker_binary_config
     }
     self.input_path = "in"
     self.output_path = "out"
     self.pc_instance_id = "test_instance_123"
     self.container_timeout = 789
     self.test_hmac_key = "CoXbp7BOEvAN9L1CB2DAORHHr3hB7wE7tpxMYm07tc0="
    def setUp(self) -> None:
        self._pc_instance = PrivateComputationInstance(
            instance_id="123",
            role=PrivateComputationRole.PARTNER,
            instances=[],
            status=PrivateComputationInstanceStatus.
            INPUT_DATA_VALIDATION_STARTED,
            status_update_ts=1600000000,
            num_pid_containers=1,
            num_mpc_containers=1,
            num_files_per_mpc_container=1,
            game_type=PrivateComputationGameType.LIFT,
            input_path=
            "https://a-test-bucket.s3.us-west-2.amazonaws.com/lift/test/input_data1.csv",
            output_dir="789",
        )

        self.onedocker_binary_config_map = defaultdict(
            lambda: OneDockerBinaryConfig(
                tmp_directory="/test_tmp_directory/",
                binary_version="latest",
                repository_path="test_path/",
            ))
コード例 #8
0
        async def _run_sub_test(wait_for_containers: bool, ) -> None:
            ip = "192.0.2.0"
            container = ContainerInstance(instance_id="123", ip_address=ip)

            mock_onedocker_svc.start_containers = MagicMock(
                return_value=[container])
            mock_onedocker_svc.wait_for_pending_containers = AsyncMock(
                return_value=[container])

            container.status = (ContainerInstanceStatus.COMPLETED
                                if wait_for_containers else
                                ContainerInstanceStatus.STARTED)
            mock_wait_for_containers_async.return_value = [container]
            test_onedocker_binary_config = OneDockerBinaryConfig(
                tmp_directory="/test_tmp_directory/",
                binary_version="latest",
                repository_path="test_path/",
            )
            stage = PIDShardStage(
                stage=UnionPIDStage.PUBLISHER_SHARD,
                instance_repository=mock_instance_repo,
                storage_svc=mock_storage_svc,
                onedocker_svc=mock_onedocker_svc,
                onedocker_binary_config=test_onedocker_binary_config,
            )
            instance_id = "444"
            stage_input = PIDStageInput(
                input_paths=["in"],
                output_paths=["out"],
                num_shards=123,
                instance_id=instance_id,
            )

            # Basic test: All good
            with patch.object(PIDShardStage, "files_exist") as mock_fe:
                mock_fe.return_value = True
                stage = PIDShardStage(
                    stage=UnionPIDStage.PUBLISHER_SHARD,
                    instance_repository=mock_instance_repo,
                    storage_svc=mock_storage_svc,
                    onedocker_svc=mock_onedocker_svc,
                    onedocker_binary_config=test_onedocker_binary_config,
                )
                status = await stage.run(
                    stage_input, wait_for_containers=wait_for_containers)
                self.assertEqual(
                    PIDStageStatus.COMPLETED
                    if wait_for_containers else PIDStageStatus.STARTED,
                    status,
                )

                mock_onedocker_svc.start_containers.assert_called_once()
                if wait_for_containers:
                    mock_wait_for_containers_async.assert_called_once()
                else:
                    mock_wait_for_containers_async.assert_not_called()
                # instance status is updated to READY, STARTED, then COMPLETED
                mock_instance_repo.read.assert_called_with(instance_id)
                self.assertEqual(mock_instance_repo.read.call_count, 4)
                self.assertEqual(mock_instance_repo.update.call_count, 4)

            # Input not ready
            with patch.object(PIDShardStage, "files_exist") as mock_fe:
                mock_fe.return_value = False
                status = await stage.run(
                    stage_input, wait_for_containers=wait_for_containers)
                self.assertEqual(PIDStageStatus.FAILED, status)

            # Multiple input paths (invariant exception)
            with patch.object(PIDShardStage, "files_exist") as mock_fe:
                with self.assertRaises(ValueError):
                    mock_fe.return_value = True
                    stage_input.input_paths = ["in1", "in2"]
                    stage = PIDShardStage(
                        stage=UnionPIDStage.PUBLISHER_SHARD,
                        instance_repository=mock_instance_repo,
                        storage_svc=mock_storage_svc,
                        onedocker_svc=mock_onedocker_svc,
                        onedocker_binary_config=test_onedocker_binary_config,
                    )
                    status = await stage.run(
                        stage_input, wait_for_containers=wait_for_containers)
                    self.assertEqual(
                        PIDStageStatus.COMPLETED
                        if wait_for_containers else PIDStageStatus.STARTED,
                        status,
                    )
コード例 #9
0
 def setUp(self) -> None:
     self.onedocker_binary_config = OneDockerBinaryConfig(
         tmp_directory="/tmp/",
         binary_version="latest",
         repository_path="test_path/",
     )
コード例 #10
0
def _build_onedocker_binary_cfg(
        onedocker_binary_config: Dict[str, Any]) -> OneDockerBinaryConfig:
    return OneDockerBinaryConfig(**onedocker_binary_config["constructor"])
コード例 #11
0
    def setUp(self) -> None:
        container_svc_patcher = patch(
            "fbpcp.service.container_aws.AWSContainerService")
        storage_svc_patcher = patch(
            "fbpcp.service.storage_s3.S3StorageService")
        mpc_instance_repo_patcher = patch(
            "fbpcs.common.repository.mpc_instance_local.LocalMPCInstanceRepository"
        )
        pid_instance_repo_patcher = patch(
            "fbpcs.pid.repository.pid_instance_local.LocalPIDInstanceRepository"
        )
        private_computation_instance_repo_patcher = patch(
            "fbpcs.private_computation.repository.private_computation_instance_local.LocalPrivateComputationInstanceRepository"
        )
        mpc_game_svc_patcher = patch("fbpcp.service.mpc_game.MPCGameService")
        container_svc = container_svc_patcher.start()
        storage_svc = storage_svc_patcher.start()
        mpc_instance_repository = mpc_instance_repo_patcher.start()
        pid_instance_repository = pid_instance_repo_patcher.start()
        private_computation_instance_repository = (
            private_computation_instance_repo_patcher.start())
        mpc_game_svc = mpc_game_svc_patcher.start()

        for patcher in (
                container_svc_patcher,
                storage_svc_patcher,
                mpc_instance_repo_patcher,
                pid_instance_repo_patcher,
                private_computation_instance_repo_patcher,
                mpc_game_svc_patcher,
        ):
            self.addCleanup(patcher.stop)

        self.onedocker_service_config = OneDockerServiceConfig(
            task_definition="test_task_definition", )

        self.onedocker_binary_config_map = defaultdict(
            lambda: OneDockerBinaryConfig(
                tmp_directory="/test_tmp_directory/",
                binary_version="latest",
                repository_path="test_path/",
            ))

        self.onedocker_service = OneDockerService(
            container_svc, self.onedocker_service_config.task_definition)

        self.mpc_service = MPCService(
            container_svc=container_svc,
            instance_repository=mpc_instance_repository,
            task_definition="test_task_definition",
            mpc_game_svc=mpc_game_svc,
        )

        self.pid_service = PIDService(
            instance_repository=pid_instance_repository,
            storage_svc=storage_svc,
            onedocker_svc=self.onedocker_service,
            onedocker_binary_config_map=self.onedocker_binary_config_map,
        )

        self.pc_validator_config = PCValidatorConfig(region="us-west-2", )

        self.private_computation_service = PrivateComputationService(
            instance_repository=private_computation_instance_repository,
            storage_svc=storage_svc,
            mpc_svc=self.mpc_service,
            pid_svc=self.pid_service,
            onedocker_svc=self.onedocker_service,
            onedocker_binary_config_map=self.onedocker_binary_config_map,
            pc_validator_config=self.pc_validator_config,
        )

        self.bolt_pcs_client = BoltPCSClient(self.private_computation_service)

        self.test_instance_id = "test_id"
        self.test_role = PrivateComputationRole.PUBLISHER
        self.test_game_type = PrivateComputationGameType.LIFT
        self.test_input_path = "in_path"
        self.test_output_path = "out_path"
        self.test_num_containers = 2
        self.test_concurrency = 1
        self.test_hmac_key = "CoXbp7BOEvAN9L1CB2DAORHHr3hB7wE7tpxMYm07tc0="