leave OneDockerBinaryConfig repository_path to default None (#1400)

Summary: Pull Request resolved: #1400 ## Why In previous Diff D35096344 (94f90b6), D35137463 (4ea6319). we added the ability to override `repository_path` for MPC/PID games. During Amazon integration, amazon found that Customized Repository Support in OneDocker only works when override config.yml file. Amazon want to take the second option in https://www.internalfb.com/intern/wiki/Private_Computation_Infra/Cloud_Infra/OneDocker_0/Customized_Repository_Support_in_OneDocker/ Define the ONEDOCKER_REPOSITORY_PATH environment variable in task definition (Thanks Ziyan) here are the code pointers that one docker svc to read the default repo path if we leave it as None * https://www.internalfb.com/code/fbsource/[ee0e5ecd17ae0dc7114aea5baecfcc352f94f03f]/fbcode/measurement/private_measurement/pcp/oss/onedocker/script/runner/onedocker_runner.py?lines=361-366 * https://www.internalfb.com/code/fbsource/[ee0e5ecd17ae0dc7114aea5baecfcc352f94f03f]/fbcode/measurement/private_measurement/pcp/oss/onedocker/script/runner/onedocker_runner.py?lines=66 detail discussion https://fb.workplace.com/groups/331044242148818/posts/538772661375974/?comment_id=539554887964418 ## What * leave oneDockerBinaryConfig repository_path to default None * not override `env_vars` if repository_path didn't specify Reviewed By: jrodal98 Differential Revision: D38392439 fbshipit-source-id: b2242f2321b44c7b6f547932739376067652df19
facebookresearch · Aug 3, 2022 · e1c367c · e1c367c
1 parent 780d961
commit e1c367c
Show file tree

Hide file tree

Showing 11 changed files with 61 additions and 23 deletions.
diff --git a/fbpcs/onedocker_binary_config.py b/fbpcs/onedocker_binary_config.py
@@ -7,17 +7,19 @@
 # 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
 @dataclass
 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
             )