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)
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)
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
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()
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
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/",
))
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,
)
def setUp(self) -> None:
self.onedocker_binary_config = OneDockerBinaryConfig(
tmp_directory="/tmp/",
binary_version="latest",
repository_path="test_path/",
)
def _build_onedocker_binary_cfg(
onedocker_binary_config: Dict[str, Any]) -> OneDockerBinaryConfig:
return OneDockerBinaryConfig(**onedocker_binary_config["constructor"])
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="