diff --git a/fbpcs/onedocker_binary_config.py b/fbpcs/onedocker_binary_config.py index b812f8138..39383ad00 100644 --- a/fbpcs/onedocker_binary_config.py +++ b/fbpcs/onedocker_binary_config.py @@ -7,10 +7,14 @@ # pyre-strict from dataclasses import dataclass +from typing import Optional from dataclasses_json import dataclass_json ONEDOCKER_REPOSITORY_PATH = "ONEDOCKER_REPOSITORY_PATH" +DEFAULT_BINARY_REPOSITORY = ( + "https://one-docker-repository-prod.s3.us-west-2.amazonaws.com/" +) @dataclass_json @@ -18,6 +22,4 @@ class OneDockerBinaryConfig: tmp_directory: str binary_version: str - repository_path: str = ( - "https://one-docker-repository-prod.s3.us-west-2.amazonaws.com/" - ) + repository_path: Optional[str] = None diff --git a/fbpcs/pc_pre_validation/constants.py b/fbpcs/pc_pre_validation/constants.py index 0f024435d..1d91bad21 100644 --- a/fbpcs/pc_pre_validation/constants.py +++ b/fbpcs/pc_pre_validation/constants.py @@ -9,6 +9,11 @@ import re from typing import Dict, List, Pattern +from fbpcs.onedocker_binary_config import ( + DEFAULT_BINARY_REPOSITORY, + ONEDOCKER_REPOSITORY_PATH, +) + from fbpcs.pc_pre_validation.binary_path import BinaryInfo INPUT_DATA_VALIDATOR_NAME = "Input Data Validator" @@ -67,9 +72,7 @@ VALID_LINE_ENDING_REGEX: Pattern[str] = re.compile(r".*(\S|\S\n)$") -DEFAULT_BINARY_REPOSITORY = ( - "https://one-docker-repository-prod.s3.us-west-2.amazonaws.com/" -) +DEFAULT_BINARY_REPOSITORY = DEFAULT_BINARY_REPOSITORY DEFAULT_BINARY_VERSION = "latest" DEFAULT_EXE_FOLDER = "/root/onedocker/package/" BINARY_INFOS: List[BinaryInfo] = [ @@ -87,5 +90,5 @@ BinaryInfo("private_attribution/shard-aggregator"), BinaryInfo("private_lift/lift"), ] -ONEDOCKER_REPOSITORY_PATH = "ONEDOCKER_REPOSITORY_PATH" +ONEDOCKER_REPOSITORY_PATH = ONEDOCKER_REPOSITORY_PATH ONEDOCKER_EXE_PATH = "ONEDOCKER_EXE_PATH" diff --git a/fbpcs/private_computation/service/pc_pre_validation_stage_service.py b/fbpcs/private_computation/service/pc_pre_validation_stage_service.py index f61deb557..9ad24e8ba 100644 --- a/fbpcs/private_computation/service/pc_pre_validation_stage_service.py +++ b/fbpcs/private_computation/service/pc_pre_validation_stage_service.py @@ -95,8 +95,10 @@ async def run_pc_pre_validation_cli( region, binary_config, ) + env_vars = {} + if binary_config.repository_path: + env_vars[ONEDOCKER_REPOSITORY_PATH] = binary_config.repository_path - env_vars = {ONEDOCKER_REPOSITORY_PATH: binary_config.repository_path} container_instances = await RunBinaryBaseService().start_containers( [cmd_args], self._onedocker_svc, diff --git a/fbpcs/private_computation/service/pid_prepare_stage_service.py b/fbpcs/private_computation/service/pid_prepare_stage_service.py index f35df7677..f3fcbafb1 100644 --- a/fbpcs/private_computation/service/pid_prepare_stage_service.py +++ b/fbpcs/private_computation/service/pid_prepare_stage_service.py @@ -125,7 +125,12 @@ async def start_pid_prepare_service( # start containers logging.info(f"{pc_role} spinning up containers") - env_vars = {ONEDOCKER_REPOSITORY_PATH: onedocker_binary_config.repository_path} + env_vars = {} + if onedocker_binary_config.repository_path: + env_vars[ + ONEDOCKER_REPOSITORY_PATH + ] = onedocker_binary_config.repository_path + pid_prepare_binary_service = PIDPrepareBinaryService() return await pid_prepare_binary_service.start_containers( cmd_args_list=args_list, diff --git a/fbpcs/private_computation/service/pid_run_protocol_stage_service.py b/fbpcs/private_computation/service/pid_run_protocol_stage_service.py index 52c0d5cbc..6d148a5af 100644 --- a/fbpcs/private_computation/service/pid_run_protocol_stage_service.py +++ b/fbpcs/private_computation/service/pid_run_protocol_stage_service.py @@ -138,7 +138,12 @@ async def start_pid_run_protocol_service( pid_protocol, pc_role ) onedocker_binary_config = self._onedocker_binary_config_map[binary_name] - env_vars = {ONEDOCKER_REPOSITORY_PATH: onedocker_binary_config.repository_path} + env_vars = {} + if onedocker_binary_config.repository_path: + env_vars[ + ONEDOCKER_REPOSITORY_PATH + ] = onedocker_binary_config.repository_path + return await pid_run_protocol_binary_service.start_containers( cmd_args_list=args_list, onedocker_svc=self._onedocker_svc, diff --git a/fbpcs/private_computation/service/pid_shard_stage_service.py b/fbpcs/private_computation/service/pid_shard_stage_service.py index e83b50739..73fb8e7f8 100644 --- a/fbpcs/private_computation/service/pid_shard_stage_service.py +++ b/fbpcs/private_computation/service/pid_shard_stage_service.py @@ -116,7 +116,12 @@ async def start_pid_shard_service( ) # start containers logging.info(f"{pc_role} spinning up containers") - env_vars = {ONEDOCKER_REPOSITORY_PATH: onedocker_binary_config.repository_path} + env_vars = {} + if onedocker_binary_config.repository_path: + env_vars[ + ONEDOCKER_REPOSITORY_PATH + ] = onedocker_binary_config.repository_path + return await sharding_binary_service.start_containers( cmd_args_list=[args], onedocker_svc=self._onedocker_svc, diff --git a/fbpcs/private_computation/service/pre_validate_service.py b/fbpcs/private_computation/service/pre_validate_service.py index f88423638..7cd8e89a7 100644 --- a/fbpcs/private_computation/service/pre_validate_service.py +++ b/fbpcs/private_computation/service/pre_validate_service.py @@ -39,7 +39,9 @@ async def run_pre_validate_async( onedocker_svc = pc_service.onedocker_svc binary_name = OneDockerBinaryNames.PC_PRE_VALIDATION.value binary_config = pc_service.onedocker_binary_config_map[binary_name] - env_vars = {ONEDOCKER_REPOSITORY_PATH: binary_config.repository_path} + env_vars = {} + if binary_config.repository_path: + env_vars[ONEDOCKER_REPOSITORY_PATH] = binary_config.repository_path cmd_args = [ get_cmd_args(input_path, region, binary_config) diff --git a/fbpcs/private_computation/service/utils.py b/fbpcs/private_computation/service/utils.py index 44c1812d3..eaa788548 100644 --- a/fbpcs/private_computation/service/utils.py +++ b/fbpcs/private_computation/service/utils.py @@ -86,10 +86,9 @@ async def create_and_start_mpc_instance( num_workers=num_containers, game_args=game_args, ) - env_vars = {} if repository_path: - env_vars["ONEDOCKER_REPOSITORY_PATH"] = repository_path + env_vars[ONEDOCKER_REPOSITORY_PATH] = repository_path return await mpc_svc.start_instance_async( instance_id=instance_id, @@ -282,7 +281,10 @@ async def start_combiner_service( multi_conversion_limit=multi_conversion_limit, log_cost=log_cost, ) - env_vars = {ONEDOCKER_REPOSITORY_PATH: binary_config.repository_path} + env_vars = {} + if binary_config.repository_path: + env_vars[ONEDOCKER_REPOSITORY_PATH] = binary_config.repository_path + return await combiner_service.start_containers( cmd_args_list=args, onedocker_svc=onedocker_svc, @@ -348,7 +350,10 @@ async def start_sharder_service( args_list.append(args_per_shard) binary_name = sharder.get_binary_name(ShardType.ROUND_ROBIN) - env_vars = {ONEDOCKER_REPOSITORY_PATH: binary_config.repository_path} + env_vars = {} + if binary_config.repository_path: + env_vars[ONEDOCKER_REPOSITORY_PATH] = binary_config.repository_path + return await sharder.start_containers( cmd_args_list=args_list, onedocker_svc=onedocker_svc, diff --git a/fbpcs/private_computation/test/service/test_pid_prepare_stage_service.py b/fbpcs/private_computation/test/service/test_pid_prepare_stage_service.py index e842acd6e..026004b2f 100644 --- a/fbpcs/private_computation/test/service/test_pid_prepare_stage_service.py +++ b/fbpcs/private_computation/test/service/test_pid_prepare_stage_service.py @@ -93,9 +93,12 @@ async def _run_sub_test( return_value=containers ) updated_pc_instance = await stage_svc.run_async(pc_instance=pc_instance) - env_vars = { - "ONEDOCKER_REPOSITORY_PATH": self.onedocker_binary_config.repository_path - } + env_vars = {} + if self.onedocker_binary_config.repository_path: + env_vars[ + "ONEDOCKER_REPOSITORY_PATH" + ] = self.onedocker_binary_config.repository_path + args_ls_expect = self.get_args_expected( pc_role, test_num_containers, max_col_cnt_expect ) diff --git a/fbpcs/private_computation/test/service/test_pid_run_protocol_stage_service.py b/fbpcs/private_computation/test/service/test_pid_run_protocol_stage_service.py index 82e9b626a..141f4508b 100644 --- a/fbpcs/private_computation/test/service/test_pid_run_protocol_stage_service.py +++ b/fbpcs/private_computation/test/service/test_pid_run_protocol_stage_service.py @@ -109,7 +109,10 @@ async def _run_sub_test( pid_protocol, pc_role ) binary_config = self.onedocker_binary_config_map[binary_name] - env_vars = {ONEDOCKER_REPOSITORY_PATH: binary_config.repository_path} + env_vars = {} + if binary_config.repository_path: + env_vars[ONEDOCKER_REPOSITORY_PATH] = binary_config.repository_path + args_str_expect = self.get_args_expect( pc_role, pid_protocol, self.use_row_numbers ) diff --git a/fbpcs/private_computation/test/service/test_pid_shard_stage_service.py b/fbpcs/private_computation/test/service/test_pid_shard_stage_service.py index af8cba906..5f5ce33fc 100644 --- a/fbpcs/private_computation/test/service/test_pid_shard_stage_service.py +++ b/fbpcs/private_computation/test/service/test_pid_shard_stage_service.py @@ -81,9 +81,12 @@ async def _run_sub_test( return_value=containers ) updated_pc_instance = await stage_svc.run_async(pc_instance=pc_instance) - env_vars = { - "ONEDOCKER_REPOSITORY_PATH": self.onedocker_binary_config.repository_path - } + env_vars = {} + if self.onedocker_binary_config.repository_path: + env_vars[ + "ONEDOCKER_REPOSITORY_PATH" + ] = self.onedocker_binary_config.repository_path + args_ls_expect = self.get_args_expect( pc_role, test_num_containers, has_hmac_key )