Repository: facebookresearch/CompilerGym

Issue:
Ship cBench-v0 by default
## 🚀 Feature

Make cBench-v0 install out of the box.

## Motivation

It is a bit irritating using the gym the first time to discover that you have to download a dataset. Better to have at least a few benchmarks immediately available.

## Pitch

Clearing obstacles out of the path of first time users is bound to be a good idea.

## Alternatives

1. Don't do it.
2. Wait for the Dataset API to be overhauled.
3. Download any dataset on demand (but still has problems - e.g. how do you know what names to use; autodownloading huge datasets; etc.)



Discussion:

Resolving patch:
diff --git a/compiler_gym/envs/compiler_env.py b/compiler_gym/envs/compiler_env.py
--- a/compiler_gym/envs/compiler_env.py
+++ b/compiler_gym/envs/compiler_env.py
@@ -401,9 +401,24 @@ def benchmark(self, benchmark: Optional[Union[str, Benchmark]]):
             warnings.warn(
                 "Changing the benchmark has no effect until reset() is called."
             )
-        if isinstance(benchmark, str) or benchmark is None:
+        if benchmark is None:
+            self.logger.debug("Unsetting the forced environment")
+            self._user_specified_benchmark_uri = None
+        elif isinstance(benchmark, str):
+            self.logger.debug("Setting benchmark by name: %s", benchmark)
+            # If the user requested a benchmark by URI, e.g.
+            # benchmark://cBench-v0/dijkstra, require the dataset (cBench-v0)
+            # automatically.
+            if self.datasets_site_path:
+                components = benchmark.split("://")
+                if len(components) == 1 or components[0] == "benchmark":
+                    components = components[-1].split("/")
+                    if len(components) > 1:
+                        self.logger.info("Requiring dataset %s", components[0])
+                        self.require_dataset(components[0])
             self._user_specified_benchmark_uri = benchmark
         elif isinstance(benchmark, Benchmark):
+            self.logger.debug("Setting benchmark data: %s", benchmark.uri)
             self._user_specified_benchmark_uri = benchmark.uri
             self._add_custom_benchmarks([benchmark])
         else:
diff --git a/compiler_gym/envs/llvm/llvm_env.py b/compiler_gym/envs/llvm/llvm_env.py
--- a/compiler_gym/envs/llvm/llvm_env.py
+++ b/compiler_gym/envs/llvm/llvm_env.py
@@ -222,6 +222,24 @@ def __init__(self, *args, **kwargs):
             },
         )
 
+    def reset(self, *args, **kwargs):
+        # The BenchmarkFactory::getBenchmark() method raises an error if there
+        # are no benchmarks to select from. Install the cBench dataset as a
+        # fallback.
+        #
+        # TODO(github.com/facebookresearch/CompilerGym/issues/45): Remove this
+        # once the dataset API has been refactored so that service-side datasets
+        # are no longer an issue.
+        try:
+            super().reset(*args, **kwargs)
+        except FileNotFoundError:
+            self.logger.warning(
+                "reset() called on servie with no benchmarks available. "
+                "Installing cBench-v0"
+            )
+            self.require_dataset("cBench-v0")
+            super().reset(*args, **kwargs)
+
     @staticmethod
     def make_benchmark(*args, **kwargs):
         """Alias to :func:`llvm.make_benchmark() <compiler_gym.envs.llvm.make_benchmark>`."""


Test patch:
diff --git a/tests/llvm/BUILD b/tests/llvm/BUILD
--- a/tests/llvm/BUILD
+++ b/tests/llvm/BUILD
@@ -88,8 +88,10 @@ py_test(
     timeout = "short",
     srcs = ["datasets_test.py"],
     deps = [
+        "//compiler_gym",
         "//compiler_gym/envs/llvm:datasets",
         "//tests:test_main",
+        "//tests/pytest_plugins:common",
     ],
 )
 
diff --git a/tests/llvm/datasets_test.py b/tests/llvm/datasets_test.py
--- a/tests/llvm/datasets_test.py
+++ b/tests/llvm/datasets_test.py
@@ -3,10 +3,16 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 """Tests for //compiler_gym/envs/llvm:datasets."""
+import os
+
+import gym
+import pytest
 
 from compiler_gym.envs.llvm import datasets
 from tests.test_main import main
 
+pytest_plugins = ["tests.pytest_plugins.common"]
+
 
 def test_validate_sha_output_okay():
     output = datasets.BenchmarkExecutionResult(
@@ -25,5 +31,56 @@ def test_validate_sha_output_invalid():
     assert datasets.validate_sha_output(output)
 
 
+def test_default_cBench_dataset_require(tmpwd, temporary_environ):
+    """Test that cBench is downloaded."""
+    del temporary_environ
+
+    os.environ["COMPILER_GYM_SITE_DATA"] = str(tmpwd / "site_data")
+    env = gym.make("llvm-v0")
+    try:
+        assert not env.benchmarks, "Sanity check"
+
+        # Datasaet is downloaded.
+        assert env.require_dataset("cBench-v0")
+        assert env.benchmarks
+
+        # Dataset is already downloaded.
+        assert not env.require_dataset("cBench-v0")
+    finally:
+        env.close()
+
+
+def test_default_cBench_on_reset(tmpwd, temporary_environ):
+    """Test that cBench is downloaded by default when no benchmarks are available."""
+    del temporary_environ
+
+    os.environ["COMPILER_GYM_SITE_DATA"] = str(tmpwd / "site_data")
+    env = gym.make("llvm-v0")
+    try:
+        assert not env.benchmarks, "Sanity check"
+
+        env.reset()
+        assert env.benchmarks
+        assert env.benchmark.startswith("benchmark://cBench-v0/")
+    finally:
+        env.close()
+
+
+@pytest.mark.parametrize("benchmark_name", ["benchmark://npb-v0/1", "npb-v0/1"])
+def test_dataset_required(tmpwd, temporary_environ, benchmark_name):
+    """Test that the required dataset is downlaoded when a benchmark is specified."""
+    del temporary_environ
+
+    os.environ["COMPILER_GYM_SITE_DATA"] = str(tmpwd / "site_data")
+    env = gym.make("llvm-v0")
+    try:
+        env.reset(benchmark=benchmark_name)
+
+        assert env.benchmarks
+        assert env.benchmark.startswith("benchmark://npb-v0/")
+    finally:
+        env.close()
+
+
 if __name__ == "__main__":
     main()
diff --git a/tests/pytest_plugins/common.py b/tests/pytest_plugins/common.py
--- a/tests/pytest_plugins/common.py
+++ b/tests/pytest_plugins/common.py
@@ -23,3 +23,14 @@ def tmpwd() -> Path:
             yield Path(d)
         finally:
             os.chdir(pwd)
+
+
+@pytest.fixture(scope="function")
+def temporary_environ():
+    """A fixture that allows you to modify os.environ without affecting other tests."""
+    old_env = os.environ.copy()
+    try:
+        yield
+    finally:
+        os.environ.clear()
+        os.environ.update(old_env)
diff --git a/tests/util/BUILD b/tests/util/BUILD
--- a/tests/util/BUILD
+++ b/tests/util/BUILD
@@ -20,6 +20,7 @@ py_test(
     deps = [
         "//compiler_gym/util",
         "//tests:test_main",
+        "//tests/pytest_plugins:common",
     ],
 )
 
diff --git a/tests/util/debug_util_test.py b/tests/util/debug_util_test.py
--- a/tests/util/debug_util_test.py
+++ b/tests/util/debug_util_test.py
@@ -6,25 +6,15 @@
 import logging
 import os
 
-import pytest
-
 from compiler_gym.util import debug_util as dbg
 from tests.test_main import main
 
-
-@pytest.fixture(scope="function")
-def temporary_environ():
-    old_env = os.environ.copy()
-    os.environ.clear()
-    try:
-        yield
-    finally:
-        os.environ.clear()
-        os.environ.update(old_env)
+pytest_plugins = ["tests.pytest_plugins.common"]
 
 
 def test_get_debug_level_environment_variable(temporary_environ):
     del temporary_environ
+    os.environ.clear()
     os.environ["COMPILER_GYM_DEBUG"] = "0"
     assert dbg.get_debug_level() == 0
     os.environ["COMPILER_GYM_DEBUG"] = "1"
@@ -33,6 +23,7 @@ def test_get_debug_level_environment_variable(temporary_environ):
 
 def test_get_and_set_debug_level(temporary_environ):
     del temporary_environ
+    os.environ.clear()
     dbg.set_debug_level(0)
     assert dbg.get_debug_level() == 0
     dbg.set_debug_level(1)
@@ -41,18 +32,21 @@ def test_get_and_set_debug_level(temporary_environ):
 
 def test_negative_debug_level(temporary_environ):
     del temporary_environ
+    os.environ.clear()
     dbg.set_debug_level(-1)
     assert dbg.get_debug_level() == 0
 
 
 def test_out_of_range_debug_level(temporary_environ):
     del temporary_environ
+    os.environ.clear()
     dbg.set_debug_level(15)
     assert dbg.get_debug_level() == 15
 
 
 def test_get_logging_level(temporary_environ):
     del temporary_environ
+    os.environ.clear()
     dbg.set_debug_level(0)
     assert dbg.get_logging_level() == logging.ERROR
     dbg.set_debug_level(1)