[datasets] Switch CompilerEnv to the new dataset API.

BUILD.bazel

@@ -14,8 +14,8 @@ exports_files([
 py_library(
     name = "CompilerGym",
     data = [
-        "//compiler_gym/third_party/cBench:benchmarks_list",
-        "//compiler_gym/third_party/cBench:crc32",
+        "//compiler_gym/third_party/cbench:benchmarks_list",
+        "//compiler_gym/third_party/cbench:crc32",
     ],
     deps = [
         "//compiler_gym",

benchmarks/BUILD

@@ -10,6 +10,7 @@ py_test(
     shard_count = 8,
     deps = [
         "//compiler_gym",
+        "//examples/example_compiler_gym_service",
         "//tests:test_main",
         "//tests/pytest_plugins:llvm",
     ],

benchmarks/bench_test.py

@@ -4,8 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 """Microbenchmarks for CompilerGym environments.
 
-To run these benchmarks within bazel, compile with optimiztions and stream the
-test output:
+To run these benchmarks an optimized build using bazel:
 
     $ bazel test -c opt --test_output=streamed //benchmarks:bench_test
 
@@ -14,91 +13,178 @@
 multiple runs using:
 
     $ pytest-benchmark compare --group-by=name --sort=fullname \
-        /tmp/compiler_gym/pytest_benchmark/*_bench_test.json
+        /tmp/compiler_gym/pytest_benchmark/*/*_bench_test.json
 """
-# pylint: disable=redefined-outer-name
 import gym
 import pytest
 
+import examples.example_compiler_gym_service as dummy
 from compiler_gym.envs import CompilerEnv, LlvmEnv, llvm
 from compiler_gym.service import CompilerGymServiceConnection
+from tests.pytest_plugins.llvm import OBSERVATION_SPACE_NAMES, REWARD_SPACE_NAMES
 from tests.test_main import main
 
-pytest_plugins = ["tests.pytest_plugins.llvm"]
 
-# Redefine this fixture since running all of the benchmarks in cBench would
-# take too long, but we do want to use at least one small and one large
-# benchmark to see both per-invocation overhead and overhead that is a result
-# of the size of the fixture.
-#
-# adpcm is small and jpeg-d is large. ghostscript is the largest but that
-# one takes too long.
 @pytest.fixture(
-    params=["cBench-v1/crc32", "cBench-v1/jpeg-d"],
-    ids=["fast_benchmark", "slow_benchmark"],
+    params=["llvm-v0", "example-cc-v0", "example-py-v0"],
+    ids=["llvm", "dummy-cc", "dummy-py"],
 )
-def benchmark_name(request) -> str:
+def env_id(request) -> str:
     yield request.param
 
 
-@pytest.fixture(params=["cBench-v1/crc32"], ids=["fast_benchmark"])
-def fast_benchmark_name(request) -> str:
-    yield request.param
-
-
-@pytest.fixture(params=["-globaldce", "-gvn"], ids=["fast_action", "slow_action"])
-def action_name(request) -> str:
+@pytest.fixture(
+    params=["llvm-v0", "example-cc-v0", "example-py-v0"],
+    ids=["llvm", "dummy-cc", "dummy-py"],
+)
+def env(request) -> CompilerEnv:
     yield request.param
 
 
-def test_make_local(benchmark):
-    benchmark(lambda: gym.make("llvm-v0").close())
-
-
-def test_make_service(benchmark):
-    service = CompilerGymServiceConnection(llvm.LLVM_SERVICE_BINARY)
+@pytest.mark.parametrize(
+    "env_id",
+    ["llvm-v0", "example-cc-v0", "example-py-v0"],
+    ids=["llvm", "dummy-cc", "dummy-py"],
+)
+def test_make_local(benchmark, env_id):
+    benchmark(lambda: gym.make(env_id).close())
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        (llvm.LLVM_SERVICE_BINARY, LlvmEnv),
+        (dummy.EXAMPLE_CC_SERVICE_BINARY, CompilerEnv),
+        (dummy.EXAMPLE_PY_SERVICE_BINARY, CompilerEnv),
+    ],
+    ids=["llvm", "dummy-cc", "dummy-py"],
+)
+def test_make_service(benchmark, args):
+    service_binary, env_class = args
+    service = CompilerGymServiceConnection(service_binary)
     try:
-        benchmark(lambda: LlvmEnv(service=service.connection.url).close())
+        benchmark(lambda: env_class(service=service.connection.url).close())
     finally:
         service.close()
 
 
-def test_reset(benchmark, env: CompilerEnv, benchmark_name):
-    env.observation_space = "Autophase"
-    benchmark(env.reset, benchmark_name)
-
-
-def test_step(benchmark, env: CompilerEnv, benchmark_name, action_name):
-    env.observation_space = "Autophase"
-    env.reward_space = "IrInstructionCount"
-
-    env.reset(benchmark_name)
-    action = env.action_space.flags.index(action_name)
-    benchmark(env.step, action)
-
-
-def test_observation(
-    benchmark, env: CompilerEnv, fast_benchmark_name, observation_space
-):
-    env.reset(fast_benchmark_name)
-    benchmark(lambda: env.observation[observation_space])
+@pytest.mark.parametrize(
+    "make_env",
+    [
+        lambda: gym.make("llvm-autophase-ic-v0", benchmark="cbench-v1/crc32"),
+        lambda: gym.make("llvm-autophase-ic-v0", benchmark="cbench-v1/jpeg-d"),
+        lambda: gym.make("example-cc-v0"),
+        lambda: gym.make("example-py-v0"),
+    ],
+    ids=["llvm;fast-benchmark", "llvm;slow-benchmark", "dummy-cc", "dummy-py"],
+)
+def test_reset(benchmark, make_env: CompilerEnv):
+    with make_env() as env:
+        benchmark(env.reset)
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        (
+            lambda: gym.make("llvm-autophase-ic-v0", benchmark="cbench-v1/crc32"),
+            "-globaldce",
+        ),
+        (lambda: gym.make("llvm-autophase-ic-v0", benchmark="cbench-v1/crc32"), "-gvn"),
+        (
+            lambda: gym.make("llvm-autophase-ic-v0", benchmark="cbench-v1/jpeg-d"),
+            "-globaldce",
+        ),
+        (
+            lambda: gym.make("llvm-autophase-ic-v0", benchmark="cbench-v1/jpeg-d"),
+            "-gvn",
+        ),
+        (lambda: gym.make("example-cc-v0"), "a"),
+        (lambda: gym.make("example-py-v0"), "a"),
+    ],
+    ids=[
+        "llvm;fast-benchmark;fast-action",
+        "llvm;fast-benchmark;slow-action",
+        "llvm;slow-benchmark;fast-action",
+        "llvm;slow-benchmark;slow-action",
+        "dummy-cc",
+        "dummy-py",
+    ],
+)
+def test_step(benchmark, args):
+    make_env, action_name = args
+    with make_env() as env:
+        env.reset()
+        action = env.action_space[action_name]
+        benchmark(env.step, action)
+
+
+_args = dict(
+    {
+        f"llvm;{obs}": (lambda: gym.make("llvm-v0", benchmark="cbench-v1/qsort"), obs)
+        for obs in OBSERVATION_SPACE_NAMES
+    },
+    **{
+        "dummy-cc": (lambda: gym.make("example-cc-v0"), "ir"),
+        "dummy-py": (lambda: gym.make("example-py-v0"), "features"),
+    },
+)
 
 
-def test_reward(benchmark, env: CompilerEnv, benchmark_name, reward_space):
-    env.reset(benchmark_name)
-    benchmark(lambda: env.reward[reward_space])
+@pytest.mark.parametrize("args", _args.values(), ids=_args.keys())
+def test_observation(benchmark, args):
+    make_env, observation_space = args
+    with make_env() as env:
+        env.reset()
+        benchmark(lambda: env.observation[observation_space])
+
+
+_args = dict(
+    {
+        f"llvm;{reward}": (
+            lambda: gym.make("llvm-v0", benchmark="cbench-v1/qsort"),
+            reward,
+        )
+        for reward in REWARD_SPACE_NAMES
+    },
+    **{
+        "dummy-cc": (lambda: gym.make("example-cc-v0"), "runtime"),
+        "dummy-py": (lambda: gym.make("example-py-v0"), "runtime"),
+    },
+)
 
 
-def test_fork(benchmark, env: CompilerEnv, benchmark_name):
-    env.reset(benchmark_name)
-    benchmark(lambda: env.fork().close())
+@pytest.mark.parametrize("args", _args.values(), ids=_args.keys())
+def test_reward(benchmark, args):
+    make_env, reward_space = args
+    with make_env() as env:
+        env.reset()
+        benchmark(lambda: env.reward[reward_space])
+
+
+@pytest.mark.parametrize(
+    "make_env",
+    [
+        lambda: gym.make("llvm-autophase-ic-v0", benchmark="cbench-v1/crc32"),
+        lambda: gym.make("llvm-autophase-ic-v0", benchmark="cbench-v1/jpeg-d"),
+        # TODO: Example service does not yet support fork() operator.
+        # lambda: gym.make("example-cc-v0"),
+        # lambda: gym.make("example-py-v0"),
+    ],
+    ids=["llvm;fast-benchmark", "llvm;slow-benchmark"],
+)
+def test_fork(benchmark, make_env):
+    with make_env() as env:
+        env.reset()
+        benchmark(lambda: env.fork().close())
 
 
 if __name__ == "__main__":
     main(
         extra_pytest_args=[
             "--benchmark-storage=/tmp/compiler_gym/pytest_benchmark",
             "--benchmark-save=bench_test",
+            "--benchmark-sort=name",
             "-x",
         ],
         debug_level=0,

compiler_gym/bin/BUILD

@@ -22,6 +22,7 @@ py_binary(
     srcs = ["datasets.py"],
     visibility = ["//visibility:public"],
     deps = [
+        ":service",
         "//compiler_gym/datasets",
         "//compiler_gym/envs",
         "//compiler_gym/util",
@@ -81,6 +82,7 @@ py_binary(
     srcs = ["service.py"],
     visibility = ["//visibility:public"],
     deps = [
+        "//compiler_gym/datasets",
         "//compiler_gym/envs",
         "//compiler_gym/spaces",
         "//compiler_gym/util",