[examples] Fix multiprocessing error during GCC init.
"""
import random
from itertools import islice, product
+from multiprocessing import Lock
from pathlib import Path
from typing import List, NamedTuple
@@ -64,6 +65,10 @@ flags.DEFINE_enum(
"objective", "obj_size", ["asm_size", "obj_size"], "Which objective to use"
)
+# Lock to prevent multiple processes all calling compiler_gym.make("gcc-v0")
+# simultaneously as this can cause issues with the docker API.
+GCC_ENV_CONSTRUCTOR_LOCK = Lock()
+
def random_search(env: CompilerEnv):
best = float("inf")
@@ -156,15 +161,20 @@ class SearchResult(NamedTuple):
def run_search(search: str, benchmark: str, seed: int) -> SearchResult:
"""Run a search and return the search class instance."""
+ with GCC_ENV_CONSTRUCTOR_LOCK:
+ env = compiler_gym.make("gcc-v0")
+
+ try:
random.seed(seed)
np.random.seed(seed)
- with compiler_gym.make("gcc-v0") as env:
env.reset(benchmark=benchmark)
env.step(env.action_space["-Os"])
baseline_size = objective(env)
env.reset(benchmark=benchmark)
best_size = _SEARCH_FUNCTIONS[search](env)
+ finally:
+ env.close()
return SearchResult(
search=search,
