[service] Add a reusable pool for backend services.

.github/workflows/ci.yaml

@@ -59,7 +59,6 @@ jobs:
                   else
                     wget https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/clang+llvm-10.0.0-x86_64-linux-gnu-ubuntu-18.04.tar.xz -O ~/llvm.tar.xz
                   fi
-                  # TODO(cummins): Remove 'v' debugging flag:
                   mkdir ~/llvm && tar xvf ~/llvm.tar.xz --strip-components 1 -C ~/llvm
                   rm ~/llvm.tar.xz
                   echo "Unpacked, testing for expected file:"

benchmarks/bench_test.py

@@ -20,9 +20,10 @@
 import gym
 import pytest
 
-import examples.example_compiler_gym_service as dummy
+import examples.example_compiler_gym_service  # noqa Environment import.
+import examples.example_compiler_gym_service as dummy  # noqa Environment import.
 from compiler_gym.envs import CompilerEnv, LlvmEnv, llvm
-from compiler_gym.service import CompilerGymServiceConnection
+from compiler_gym.service import CompilerGymServiceConnection, ConnectionOpts
 from compiler_gym.service.client_service_compiler_env import ClientServiceCompilerEnv
 from tests.pytest_plugins.llvm import OBSERVATION_SPACE_NAMES, REWARD_SPACE_NAMES
 from tests.test_main import main
@@ -46,10 +47,10 @@ def env(request) -> CompilerEnv:
 
 @pytest.mark.parametrize(
     "env_id",
-    ["llvm-v0", "example-cc-v0", "example-py-v0"],
-    ids=["llvm", "dummy-cc", "dummy-py"],
+    ["llvm-v0", "example-cc-v0", "example-py-v0", "loop_tool-v0"],
+    ids=["llvm", "dummy-cc", "dummy-py", "loop_tool"],
 )
-def test_make_local(benchmark, env_id):
+def test_make_env(benchmark, env_id):
     benchmark(lambda: gym.make(env_id).close())
 
 
@@ -64,7 +65,7 @@ def test_make_local(benchmark, env_id):
 )
 def test_make_service(benchmark, args):
     service_binary, env_class = args
-    service = CompilerGymServiceConnection(service_binary)
+    service = CompilerGymServiceConnection(service_binary, ConnectionOpts())
     try:
         benchmark(lambda: env_class(service=service.connection.url).close())
     finally:

compiler_gym/bin/service.py

@@ -266,7 +266,9 @@ def main(argv):
 
     if FLAGS.run_on_port:
         assert FLAGS.env, "Must specify an --env to run"
-        settings = ConnectionOpts(script_args=["--port", str(FLAGS.run_on_port)])
+        settings = ConnectionOpts(
+            script_args=frozenset(["--port", str(FLAGS.run_on_port)])
+        )
         with gym.make(FLAGS.env, connection_settings=settings) as env:
             print(
                 f"=== Started a service on port {FLAGS.run_on_port}. Use C-c to terminate. ==="

compiler_gym/envs/gcc/__init__.py

@@ -16,7 +16,7 @@
 
 register(
     id="gcc-v0",
-    entry_point="compiler_gym.envs.gcc:GccEnv",
+    entry_point="compiler_gym.envs.gcc.gcc_env:make",
     kwargs={"service": GCC_SERVICE_BINARY},
 )
 

compiler_gym/envs/gcc/gcc_env.py

@@ -7,6 +7,7 @@
 import json
 import pickle
 from pathlib import Path
+from threading import Lock
 from typing import Any, Dict, List, Optional, Union
 
 from compiler_gym.datasets import Benchmark
@@ -15,6 +16,7 @@
 from compiler_gym.envs.gcc.gcc_rewards import AsmSizeReward, ObjSizeReward
 from compiler_gym.service import ConnectionOpts
 from compiler_gym.service.client_service_compiler_env import ClientServiceCompilerEnv
+from compiler_gym.service.connection_pool import ServiceConnectionPoolBase
 from compiler_gym.spaces import Reward
 from compiler_gym.util.decorators import memoized_property
 from compiler_gym.util.gym_type_hints import ObservationType, OptionalArgumentValue
@@ -63,9 +65,11 @@ def __init__(
 
         :raises ServiceInitError: If the requested GCC version cannot be used.
         """
-        connection_settings = connection_settings or ConnectionOpts()
         # Pass the executable path via an environment variable
-        connection_settings.script_env = {"CC": gcc_bin}
+        connection_settings = connection_settings or ConnectionOpts()
+        connection_settings.script_env = connection_settings.script_env.set(
+            "CC", gcc_bin
+        )
 
         # Eagerly create a GCC compiler instance now because:
         #
@@ -76,6 +80,13 @@ def __init__(
         #    initialization may time out.
         Gcc(bin=gcc_bin)
 
+        # NOTE(github.com/facebookresearch/CompilerGym/pull/583): The GCC
+        # environment stalls on the StartSession() RPC call when service
+        # connection caching is enabled. I believe this has something to do with
+        # the runtime code generation, but have not been able to diagnose it
+        # yet. For now, disable service connection caching for GCC environments.
+        kwargs["service_pool"] = ServiceConnectionPoolBase()
+
         super().__init__(
             *args,
             **kwargs,
@@ -88,6 +99,9 @@ def __init__(
         )
         self._timeout = timeout
 
+    def commandline_to_actions(self, commandline: str) -> List[int]:
+        return NotImplementedError
+
     def reset(
         self,
         benchmark: Optional[Union[str, Benchmark]] = None,
@@ -213,3 +227,19 @@ def _init_kwargs(self) -> Dict[str, Any]:
             "gcc_bin": self.gcc_spec.gcc.bin,
             **super()._init_kwargs(),
         }
+
+
+_GCC_ENV_DOCKER_CONSTRUCTOR_LOCK = Lock()
+
+
+def make(*args, gcc_bin: Union[str, Path] = DEFAULT_GCC, **kwargs):
+    """Construct a GccEnv class using a lock to ensure thread exclusivity.
+
+    This is to prevent multiple threads running the docker initialization
+    routines simultaneously as this can cause issues with the docker API.
+    """
+    if gcc_bin.startswith("docker:"):
+        with _GCC_ENV_DOCKER_CONSTRUCTOR_LOCK:
+            return GccEnv(*args, gcc_bin=gcc_bin, **kwargs)
+    else:
+        return GccEnv(*args, gcc_bin=gcc_bin, **kwargs)

compiler_gym/envs/llvm/datasets/cbench.py

@@ -234,6 +234,10 @@ def download_cBench_runtime_data() -> bool:
     if (cbench_data / "unpacked").is_file():
         return False
     else:
+        logger.warning(
+            "Installing the cBench runtime inputs. This may take a few moments ..."
+        )
+
         # Clean up any partially-extracted data directory.
         if cbench_data.is_dir():
             shutil.rmtree(cbench_data)

compiler_gym/requirements.txt

@@ -2,6 +2,7 @@ absl-py>=0.10.0
 deprecated>=1.2.12
 docker>=4.0.0
 fasteners>=0.15
+frozendict>=1.0.0
 grpcio>=1.32.0,<1.44.0
 gym>=0.18.0,<0.21
 humanize>=2.6.0

compiler_gym/service/BUILD

@@ -9,14 +9,15 @@ py_library(
     name = "service",
     srcs = [
         "__init__.py",
+        "connection_pool.py",
     ],
     visibility = ["//visibility:public"],
     deps = [
-        ":compilation_session",
-        ":connection",
         # TODO(github.com/facebookresearch/CompilerGym/pull/633):
-        # add this after circular dependencies are resolved
+        # add this after circular dependencies are resolved:
         # ":client_service_compiler_env",
+        ":compilation_session",
+        ":connection",
         ":service_cache",
         "//compiler_gym/errors",
         "//compiler_gym/service/proto",

compiler_gym/service/CMakeLists.txt

@@ -10,12 +10,13 @@ cg_py_library(
     service
   SRCS
     "__init__.py"
+    "connection_pool.py"
   DEPS
+    # TODO(github.com/facebookresearch/CompilerGym/pull/633):
+    # add this after circular dependencies are resolved:
+    # ::client_service_compiler_env
     ::compilation_session
     ::connection
-    # TODO(github.com/facebookresearch/CompilerGym/pull/633):
-    # add this after circular dependencies are resolved
-    #::client_service_compiler_env
     ::service_cache
     compiler_gym::errors::errors
     compiler_gym::service::proto::proto

compiler_gym/service/__init__.py

@@ -14,12 +14,14 @@
     ServiceTransportError,
     SessionNotFound,
 )
+from compiler_gym.service.connection_pool import ServiceConnectionPool
 
 __all__ = [
-    "CompilerGymServiceConnection",
     "CompilationSession",
+    "CompilerGymServiceConnection",
     "ConnectionOpts",
     "EnvironmentNotSupported",
+    "ServiceConnectionPool",
     "ServiceError",
     "ServiceInitError",
     "ServiceIsClosed",