We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Documentation for compiler_gym.wrappers is incomplete
## 📚 Documentation

The module docstring for `compiler_gym/wrappers/__init__.py` simply reads:

> The `compiler_gym.wrappers` module provides.

👎 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `compiler_gym/wrappers/__init__.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates.
2 #
3 # This source code is licensed under the MIT license found in the
4 # LICENSE file in the root directory of this source tree.
5 """The :code:`compiler_gym.wrappers` module provides.
6 """
7 from compiler_gym.wrappers.commandline import (
8     CommandlineWithTerminalAction,
9     ConstrainedCommandline,
10 )
11 from compiler_gym.wrappers.core import (
12     ActionWrapper,
13     CompilerEnvWrapper,
14     ObservationWrapper,
15     RewardWrapper,
16 )
17 from compiler_gym.wrappers.datasets import (
18     CycleOverBenchmarks,
19     CycleOverBenchmarksIterator,
20     IterateOverBenchmarks,
21     RandomOrderBenchmarks,
22 )
23 from compiler_gym.wrappers.llvm import RuntimePointEstimateReward
24 from compiler_gym.wrappers.time_limit import TimeLimit
25 
26 __all__ = [
27     "ActionWrapper",
28     "CommandlineWithTerminalAction",
29     "CompilerEnvWrapper",
30     "ConstrainedCommandline",
31     "CycleOverBenchmarks",
32     "CycleOverBenchmarksIterator",
33     "IterateOverBenchmarks",
34     "ObservationWrapper",
35     "RandomOrderBenchmarks",
36     "RewardWrapper",
37     "RuntimePointEstimateReward",
38     "TimeLimit",
39 ]
40
```

Path: `compiler_gym/wrappers/datasets.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates.
2 #
3 # This source code is licensed under the MIT license found in the
4 # LICENSE file in the root directory of this source tree.
5 from itertools import cycle
6 from typing import Callable, Iterable, Optional, Union
7 
8 import numpy as np
9 
10 from compiler_gym.datasets import Benchmark
11 from compiler_gym.envs import CompilerEnv
12 from compiler_gym.util.parallelization import thread_safe_tee
13 from compiler_gym.wrappers.core import CompilerEnvWrapper
14 
15 BenchmarkLike = Union[str, Benchmark]
16 
17 
18 class IterateOverBenchmarks(CompilerEnvWrapper):
19     """Iterate over a (possibly infinite) sequence of benchmarks on each call to
20     reset(). Will raise :code:`StopIteration` on :meth:`reset()
21     <compiler_gym.envs.CompilerEnv.reset>` once the iterator is exhausted. Use
22     :class:`CycleOverBenchmarks` or :class:`RandomOrderBenchmarks` for wrappers
23     which will loop over the benchmarks.
24     """
25 
26     def __init__(
27         self,
28         env: CompilerEnv,
29         benchmarks: Iterable[BenchmarkLike],
30         fork_shares_iterator: bool = False,
31     ):
32         """Constructor.
33 
34         :param env: The environment to wrap.
35 
36         :param benchmarks: An iterable sequence of benchmarks.
37 
38         :param fork_shares_iterator: If :code:`True`, the :code:`benchmarks`
39             iterator will bet shared by a forked environment created by
40             :meth:`env.fork() <compiler_gym.envs.CompilerEnv.fork>`. This means
41             that calling :meth:`env.reset()
42             <compiler_gym.envs.CompilerEnv.reset>` with one environment will
43             advance the iterator in the other. If :code:`False`, forked
44             environments will use :code:`itertools.tee()` to create a copy of
45             the iterator so that each iterator may advance independently.
46             However, this requires shared buffers between the environments which
47             can lead to memory overheads if :meth:`env.reset()
48             <compiler_gym.envs.CompilerEnv.reset>` is called many times more in
49             one environment than the other.
50         """
51         super().__init__(env)
52         self.benchmarks = iter(benchmarks)
53         self.fork_shares_iterator = fork_shares_iterator
54 
55     def reset(self, benchmark: Optional[BenchmarkLike] = None, **kwargs):
56         if benchmark is not None:
57             raise TypeError("Benchmark passed to IterateOverBenchmarks.reset()")
58         benchmark: BenchmarkLike = next(self.benchmarks)
59         return self.env.reset(benchmark=benchmark)
60 
61     def fork(self) -> "IterateOverBenchmarks":
62         if self.fork_shares_iterator:
63             other_benchmarks_iterator = self.benchmarks
64         else:
65             self.benchmarks, other_benchmarks_iterator = thread_safe_tee(
66                 self.benchmarks
67             )
68         return IterateOverBenchmarks(
69             env=self.env.fork(),
70             benchmarks=other_benchmarks_iterator,
71             fork_shares_iterator=self.fork_shares_iterator,
72         )
73 
74 
75 class CycleOverBenchmarks(IterateOverBenchmarks):
76     """Cycle through a list of benchmarks on each call to :meth:`reset()
77     <compiler_gym.envs.CompilerEnv.reset>`. Same as
78     :class:`IterateOverBenchmarks` except the list of benchmarks repeats once
79     exhausted.
80     """
81 
82     def __init__(
83         self,
84         env: CompilerEnv,
85         benchmarks: Iterable[BenchmarkLike],
86         fork_shares_iterator: bool = False,
87     ):
88         """Constructor.
89 
90         :param env: The environment to wrap.
91 
92         :param benchmarks: An iterable sequence of benchmarks.
93 
94         :param fork_shares_iterator: If :code:`True`, the :code:`benchmarks`
95             iterator will be shared by a forked environment created by
96             :meth:`env.fork() <compiler_gym.envs.CompilerEnv.fork>`. This means
97             that calling :meth:`env.reset()
98             <compiler_gym.envs.CompilerEnv.reset>` with one environment will
99             advance the iterator in the other. If :code:`False`, forked
100             environments will use :code:`itertools.tee()` to create a copy of
101             the iterator so that each iterator may advance independently.
102             However, this requires shared buffers between the environments which
103             can lead to memory overheads if :meth:`env.reset()
104             <compiler_gym.envs.CompilerEnv.reset>` is called many times more in
105             one environment than the other.
106         """
107         super().__init__(
108             env, benchmarks=cycle(benchmarks), fork_shares_iterator=fork_shares_iterator
109         )
110 
111 
112 class CycleOverBenchmarksIterator(CompilerEnvWrapper):
113     """Same as :class:`CycleOverBenchmarks
114     <compiler_gym.wrappers.CycleOverBenchmarks>` except that the user generates
115     the iterator.
116     """
117 
118     def __init__(
119         self,
120         env: CompilerEnv,
121         make_benchmark_iterator: Callable[[], Iterable[BenchmarkLike]],
122     ):
123         """Constructor.
124 
125         :param env: The environment to wrap.
126 
127         :param make_benchmark_iterator: A callback that returns an iterator over
128             a sequence of benchmarks. Once the iterator is exhausted, this
129             callback is called to produce a new iterator.
130         """
131         super().__init__(env)
132         self.make_benchmark_iterator = make_benchmark_iterator
133         self.benchmarks = iter(self.make_benchmark_iterator())
134 
135     def reset(self, benchmark: Optional[BenchmarkLike] = None, **kwargs):
136         if benchmark is not None:
137             raise TypeError("Benchmark passed toIterateOverBenchmarks.reset()")
138         try:
139             benchmark: BenchmarkLike = next(self.benchmarks)
140         except StopIteration:
141             self.benchmarks = iter(self.make_benchmark_iterator())
142             benchmark: BenchmarkLike = next(self.benchmarks)
143 
144         return self.env.reset(benchmark=benchmark)
145 
146     def fork(self) -> "CycleOverBenchmarksIterator":
147         return CycleOverBenchmarksIterator(
148             env=self.env.fork(),
149             make_benchmark_iterator=self.make_benchmark_iterator,
150         )
151 
152 
153 class RandomOrderBenchmarks(IterateOverBenchmarks):
154     """Select randomly from a list of benchmarks on each call to :meth:`reset()
155     <compiler_gym.envs.CompilerEnv.reset>`.
156 
157     .. note::
158 
159         Uniform random selection is provided by evaluating the input benchmarks
160         iterator into a list and sampling randomly from the list. This will not
161         work for random iteration over infinite or very large iterables of
162         benchmarks.
163     """
164 
165     def __init__(
166         self,
167         env: CompilerEnv,
168         benchmarks: Iterable[BenchmarkLike],
169         rng: Optional[np.random.Generator] = None,
170     ):
171         """Constructor.
172 
173         :param env: The environment to wrap.
174 
175         :param benchmarks: An iterable sequence of benchmarks. The entirety of
176             this input iterator is evaluated during construction.
177 
178         :param rng: A random number generator to use for random benchmark
179             selection.
180         """
181         self._all_benchmarks = list(benchmarks)
182         rng = rng or np.random.default_rng()
183         super().__init__(
184             env,
185             benchmarks=(rng.choice(self._all_benchmarks) for _ in iter(int, 1)),
186             fork_shares_iterator=True,
187         )
188 
189     def fork(self) -> "IterateOverBenchmarks":
190         """Fork the random order benchmark wrapper.
191 
192         Note that RNG state is not copied to forked environments.
193         """
194         return IterateOverBenchmarks(
195             env=self.env.fork(), benchmarks=self._all_benchmarks
196         )
197
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/compiler_gym/wrappers/__init__.py b/compiler_gym/wrappers/__init__.py
--- a/compiler_gym/wrappers/__init__.py
+++ b/compiler_gym/wrappers/__init__.py
@@ -2,7 +2,30 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-"""The :code:`compiler_gym.wrappers` module provides.
+"""The :code:`compiler_gym.wrappers` module provides a set of classes that can
+be used to transform an environment in a modular way.
+
+For example:
+
+    >>> env = compiler_gym.make("llvm-v0")
+    >>> env = TimeLimit(env, n=10)
+    >>> env = CycleOverBenchmarks(
+    ...     env,
+    ...     benchmarks=[
+    ...         "benchmark://cbench-v1/crc32",
+    ...         "benchmark://cbench-v1/qsort",
+    ...     ],
+    ... )
+
+.. warning::
+
+    CompilerGym environments are incompatible with the `OpenAI Gym wrappers
+    <https://github.com/openai/gym/tree/master/gym/wrappers>`_. This is because
+    CompilerGym extends the environment API with additional arguments and
+    methods. You must use the wrappers from this module when wrapping
+    CompilerGym environments. We provide a set of base wrappers that are
+    equivalent to those in OpenAI Gym that you can use to write your own
+    wrappers.
 """
 from compiler_gym.wrappers.commandline import (
     CommandlineWithTerminalAction,
diff --git a/compiler_gym/wrappers/datasets.py b/compiler_gym/wrappers/datasets.py
--- a/compiler_gym/wrappers/datasets.py
+++ b/compiler_gym/wrappers/datasets.py
@@ -157,9 +157,11 @@
     .. note::
 
         Uniform random selection is provided by evaluating the input benchmarks
-        iterator into a list and sampling randomly from the list. This will not
-        work for random iteration over infinite or very large iterables of
-        benchmarks.
+        iterator into a list and sampling randomly from the list. For very large
+        and infinite iterables of benchmarks you must use the
+        :class:`IterateOverBenchmarks
+        <compiler_gym.wrappers.IterateOverBenchmarks>` wrapper with your own
+        random sampling iterator.
     """
 
     def __init__(

{"golden_diff": "diff --git a/compiler_gym/wrappers/__init__.py b/compiler_gym/wrappers/__init__.py\n--- a/compiler_gym/wrappers/__init__.py\n+++ b/compiler_gym/wrappers/__init__.py\n@@ -2,7 +2,30 @@\n #\n # This source code is licensed under the MIT license found in the\n # LICENSE file in the root directory of this source tree.\n-\"\"\"The :code:`compiler_gym.wrappers` module provides.\n+\"\"\"The :code:`compiler_gym.wrappers` module provides a set of classes that can\n+be used to transform an environment in a modular way.\n+\n+For example:\n+\n+    >>> env = compiler_gym.make(\"llvm-v0\")\n+    >>> env = TimeLimit(env, n=10)\n+    >>> env = CycleOverBenchmarks(\n+    ...     env,\n+    ...     benchmarks=[\n+    ...         \"benchmark://cbench-v1/crc32\",\n+    ...         \"benchmark://cbench-v1/qsort\",\n+    ...     ],\n+    ... )\n+\n+.. warning::\n+\n+    CompilerGym environments are incompatible with the `OpenAI Gym wrappers\n+    <https://github.com/openai/gym/tree/master/gym/wrappers>`_. This is because\n+    CompilerGym extends the environment API with additional arguments and\n+    methods. You must use the wrappers from this module when wrapping\n+    CompilerGym environments. We provide a set of base wrappers that are\n+    equivalent to those in OpenAI Gym that you can use to write your own\n+    wrappers.\n \"\"\"\n from compiler_gym.wrappers.commandline import (\n     CommandlineWithTerminalAction,\ndiff --git a/compiler_gym/wrappers/datasets.py b/compiler_gym/wrappers/datasets.py\n--- a/compiler_gym/wrappers/datasets.py\n+++ b/compiler_gym/wrappers/datasets.py\n@@ -157,9 +157,11 @@\n     .. note::\n \n         Uniform random selection is provided by evaluating the input benchmarks\n-        iterator into a list and sampling randomly from the list. This will not\n-        work for random iteration over infinite or very large iterables of\n-        benchmarks.\n+        iterator into a list and sampling randomly from the list. For very large\n+        and infinite iterables of benchmarks you must use the\n+        :class:`IterateOverBenchmarks\n+        <compiler_gym.wrappers.IterateOverBenchmarks>` wrapper with your own\n+        random sampling iterator.\n     \"\"\"\n \n     def __init__(\n", "issue": "Documentation for compiler_gym.wrappers is incomplete\n## \ud83d\udcda Documentation\r\n\r\nThe module docstring for `compiler_gym/wrappers/__init__.py` simply reads:\r\n\r\n> The `compiler_gym.wrappers` module provides.\r\n\r\n\ud83d\udc4e \n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"The :code:`compiler_gym.wrappers` module provides.\n\"\"\"\nfrom compiler_gym.wrappers.commandline import (\n    CommandlineWithTerminalAction,\n    ConstrainedCommandline,\n)\nfrom compiler_gym.wrappers.core import (\n    ActionWrapper,\n    CompilerEnvWrapper,\n    ObservationWrapper,\n    RewardWrapper,\n)\nfrom compiler_gym.wrappers.datasets import (\n    CycleOverBenchmarks,\n    CycleOverBenchmarksIterator,\n    IterateOverBenchmarks,\n    RandomOrderBenchmarks,\n)\nfrom compiler_gym.wrappers.llvm import RuntimePointEstimateReward\nfrom compiler_gym.wrappers.time_limit import TimeLimit\n\n__all__ = [\n    \"ActionWrapper\",\n    \"CommandlineWithTerminalAction\",\n    \"CompilerEnvWrapper\",\n    \"ConstrainedCommandline\",\n    \"CycleOverBenchmarks\",\n    \"CycleOverBenchmarksIterator\",\n    \"IterateOverBenchmarks\",\n    \"ObservationWrapper\",\n    \"RandomOrderBenchmarks\",\n    \"RewardWrapper\",\n    \"RuntimePointEstimateReward\",\n    \"TimeLimit\",\n]\n", "path": "compiler_gym/wrappers/__init__.py"}, {"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\nfrom itertools import cycle\nfrom typing import Callable, Iterable, Optional, Union\n\nimport numpy as np\n\nfrom compiler_gym.datasets import Benchmark\nfrom compiler_gym.envs import CompilerEnv\nfrom compiler_gym.util.parallelization import thread_safe_tee\nfrom compiler_gym.wrappers.core import CompilerEnvWrapper\n\nBenchmarkLike = Union[str, Benchmark]\n\n\nclass IterateOverBenchmarks(CompilerEnvWrapper):\n    \"\"\"Iterate over a (possibly infinite) sequence of benchmarks on each call to\n    reset(). Will raise :code:`StopIteration` on :meth:`reset()\n    <compiler_gym.envs.CompilerEnv.reset>` once the iterator is exhausted. Use\n    :class:`CycleOverBenchmarks` or :class:`RandomOrderBenchmarks` for wrappers\n    which will loop over the benchmarks.\n    \"\"\"\n\n    def __init__(\n        self,\n        env: CompilerEnv,\n        benchmarks: Iterable[BenchmarkLike],\n        fork_shares_iterator: bool = False,\n    ):\n        \"\"\"Constructor.\n\n        :param env: The environment to wrap.\n\n        :param benchmarks: An iterable sequence of benchmarks.\n\n        :param fork_shares_iterator: If :code:`True`, the :code:`benchmarks`\n            iterator will bet shared by a forked environment created by\n            :meth:`env.fork() <compiler_gym.envs.CompilerEnv.fork>`. This means\n            that calling :meth:`env.reset()\n            <compiler_gym.envs.CompilerEnv.reset>` with one environment will\n            advance the iterator in the other. If :code:`False`, forked\n            environments will use :code:`itertools.tee()` to create a copy of\n            the iterator so that each iterator may advance independently.\n            However, this requires shared buffers between the environments which\n            can lead to memory overheads if :meth:`env.reset()\n            <compiler_gym.envs.CompilerEnv.reset>` is called many times more in\n            one environment than the other.\n        \"\"\"\n        super().__init__(env)\n        self.benchmarks = iter(benchmarks)\n        self.fork_shares_iterator = fork_shares_iterator\n\n    def reset(self, benchmark: Optional[BenchmarkLike] = None, **kwargs):\n        if benchmark is not None:\n            raise TypeError(\"Benchmark passed to IterateOverBenchmarks.reset()\")\n        benchmark: BenchmarkLike = next(self.benchmarks)\n        return self.env.reset(benchmark=benchmark)\n\n    def fork(self) -> \"IterateOverBenchmarks\":\n        if self.fork_shares_iterator:\n            other_benchmarks_iterator = self.benchmarks\n        else:\n            self.benchmarks, other_benchmarks_iterator = thread_safe_tee(\n                self.benchmarks\n            )\n        return IterateOverBenchmarks(\n            env=self.env.fork(),\n            benchmarks=other_benchmarks_iterator,\n            fork_shares_iterator=self.fork_shares_iterator,\n        )\n\n\nclass CycleOverBenchmarks(IterateOverBenchmarks):\n    \"\"\"Cycle through a list of benchmarks on each call to :meth:`reset()\n    <compiler_gym.envs.CompilerEnv.reset>`. Same as\n    :class:`IterateOverBenchmarks` except the list of benchmarks repeats once\n    exhausted.\n    \"\"\"\n\n    def __init__(\n        self,\n        env: CompilerEnv,\n        benchmarks: Iterable[BenchmarkLike],\n        fork_shares_iterator: bool = False,\n    ):\n        \"\"\"Constructor.\n\n        :param env: The environment to wrap.\n\n        :param benchmarks: An iterable sequence of benchmarks.\n\n        :param fork_shares_iterator: If :code:`True`, the :code:`benchmarks`\n            iterator will be shared by a forked environment created by\n            :meth:`env.fork() <compiler_gym.envs.CompilerEnv.fork>`. This means\n            that calling :meth:`env.reset()\n            <compiler_gym.envs.CompilerEnv.reset>` with one environment will\n            advance the iterator in the other. If :code:`False`, forked\n            environments will use :code:`itertools.tee()` to create a copy of\n            the iterator so that each iterator may advance independently.\n            However, this requires shared buffers between the environments which\n            can lead to memory overheads if :meth:`env.reset()\n            <compiler_gym.envs.CompilerEnv.reset>` is called many times more in\n            one environment than the other.\n        \"\"\"\n        super().__init__(\n            env, benchmarks=cycle(benchmarks), fork_shares_iterator=fork_shares_iterator\n        )\n\n\nclass CycleOverBenchmarksIterator(CompilerEnvWrapper):\n    \"\"\"Same as :class:`CycleOverBenchmarks\n    <compiler_gym.wrappers.CycleOverBenchmarks>` except that the user generates\n    the iterator.\n    \"\"\"\n\n    def __init__(\n        self,\n        env: CompilerEnv,\n        make_benchmark_iterator: Callable[[], Iterable[BenchmarkLike]],\n    ):\n        \"\"\"Constructor.\n\n        :param env: The environment to wrap.\n\n        :param make_benchmark_iterator: A callback that returns an iterator over\n            a sequence of benchmarks. Once the iterator is exhausted, this\n            callback is called to produce a new iterator.\n        \"\"\"\n        super().__init__(env)\n        self.make_benchmark_iterator = make_benchmark_iterator\n        self.benchmarks = iter(self.make_benchmark_iterator())\n\n    def reset(self, benchmark: Optional[BenchmarkLike] = None, **kwargs):\n        if benchmark is not None:\n            raise TypeError(\"Benchmark passed toIterateOverBenchmarks.reset()\")\n        try:\n            benchmark: BenchmarkLike = next(self.benchmarks)\n        except StopIteration:\n            self.benchmarks = iter(self.make_benchmark_iterator())\n            benchmark: BenchmarkLike = next(self.benchmarks)\n\n        return self.env.reset(benchmark=benchmark)\n\n    def fork(self) -> \"CycleOverBenchmarksIterator\":\n        return CycleOverBenchmarksIterator(\n            env=self.env.fork(),\n            make_benchmark_iterator=self.make_benchmark_iterator,\n        )\n\n\nclass RandomOrderBenchmarks(IterateOverBenchmarks):\n    \"\"\"Select randomly from a list of benchmarks on each call to :meth:`reset()\n    <compiler_gym.envs.CompilerEnv.reset>`.\n\n    .. note::\n\n        Uniform random selection is provided by evaluating the input benchmarks\n        iterator into a list and sampling randomly from the list. This will not\n        work for random iteration over infinite or very large iterables of\n        benchmarks.\n    \"\"\"\n\n    def __init__(\n        self,\n        env: CompilerEnv,\n        benchmarks: Iterable[BenchmarkLike],\n        rng: Optional[np.random.Generator] = None,\n    ):\n        \"\"\"Constructor.\n\n        :param env: The environment to wrap.\n\n        :param benchmarks: An iterable sequence of benchmarks. The entirety of\n            this input iterator is evaluated during construction.\n\n        :param rng: A random number generator to use for random benchmark\n            selection.\n        \"\"\"\n        self._all_benchmarks = list(benchmarks)\n        rng = rng or np.random.default_rng()\n        super().__init__(\n            env,\n            benchmarks=(rng.choice(self._all_benchmarks) for _ in iter(int, 1)),\n            fork_shares_iterator=True,\n        )\n\n    def fork(self) -> \"IterateOverBenchmarks\":\n        \"\"\"Fork the random order benchmark wrapper.\n\n        Note that RNG state is not copied to forked environments.\n        \"\"\"\n        return IterateOverBenchmarks(\n            env=self.env.fork(), benchmarks=self._all_benchmarks\n        )\n", "path": "compiler_gym/wrappers/datasets.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"The :code:`compiler_gym.wrappers` module provides a set of classes that can\nbe used to transform an environment in a modular way.\n\nFor example:\n\n    >>> env = compiler_gym.make(\"llvm-v0\")\n    >>> env = TimeLimit(env, n=10)\n    >>> env = CycleOverBenchmarks(\n    ...     env,\n    ...     benchmarks=[\n    ...         \"benchmark://cbench-v1/crc32\",\n    ...         \"benchmark://cbench-v1/qsort\",\n    ...     ],\n    ... )\n\n.. warning::\n\n    CompilerGym environments are incompatible with the `OpenAI Gym wrappers\n    <https://github.com/openai/gym/tree/master/gym/wrappers>`_. This is because\n    CompilerGym extends the environment API with additional arguments and\n    methods. You must use the wrappers from this module when wrapping\n    CompilerGym environments. We provide a set of base wrappers that are\n    equivalent to those in OpenAI Gym that you can use to write your own\n    wrappers.\n\"\"\"\nfrom compiler_gym.wrappers.commandline import (\n    CommandlineWithTerminalAction,\n    ConstrainedCommandline,\n)\nfrom compiler_gym.wrappers.core import (\n    ActionWrapper,\n    CompilerEnvWrapper,\n    ObservationWrapper,\n    RewardWrapper,\n)\nfrom compiler_gym.wrappers.datasets import (\n    CycleOverBenchmarks,\n    CycleOverBenchmarksIterator,\n    IterateOverBenchmarks,\n    RandomOrderBenchmarks,\n)\nfrom compiler_gym.wrappers.llvm import RuntimePointEstimateReward\nfrom compiler_gym.wrappers.time_limit import TimeLimit\n\n__all__ = [\n    \"ActionWrapper\",\n    \"CommandlineWithTerminalAction\",\n    \"CompilerEnvWrapper\",\n    \"ConstrainedCommandline\",\n    \"CycleOverBenchmarks\",\n    \"CycleOverBenchmarksIterator\",\n    \"IterateOverBenchmarks\",\n    \"ObservationWrapper\",\n    \"RandomOrderBenchmarks\",\n    \"RewardWrapper\",\n    \"RuntimePointEstimateReward\",\n    \"TimeLimit\",\n]\n", "path": "compiler_gym/wrappers/__init__.py"}, {"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\nfrom itertools import cycle\nfrom typing import Callable, Iterable, Optional, Union\n\nimport numpy as np\n\nfrom compiler_gym.datasets import Benchmark\nfrom compiler_gym.envs import CompilerEnv\nfrom compiler_gym.util.parallelization import thread_safe_tee\nfrom compiler_gym.wrappers.core import CompilerEnvWrapper\n\nBenchmarkLike = Union[str, Benchmark]\n\n\nclass IterateOverBenchmarks(CompilerEnvWrapper):\n    \"\"\"Iterate over a (possibly infinite) sequence of benchmarks on each call to\n    reset(). Will raise :code:`StopIteration` on :meth:`reset()\n    <compiler_gym.envs.CompilerEnv.reset>` once the iterator is exhausted. Use\n    :class:`CycleOverBenchmarks` or :class:`RandomOrderBenchmarks` for wrappers\n    which will loop over the benchmarks.\n    \"\"\"\n\n    def __init__(\n        self,\n        env: CompilerEnv,\n        benchmarks: Iterable[BenchmarkLike],\n        fork_shares_iterator: bool = False,\n    ):\n        \"\"\"Constructor.\n\n        :param env: The environment to wrap.\n\n        :param benchmarks: An iterable sequence of benchmarks.\n\n        :param fork_shares_iterator: If :code:`True`, the :code:`benchmarks`\n            iterator will bet shared by a forked environment created by\n            :meth:`env.fork() <compiler_gym.envs.CompilerEnv.fork>`. This means\n            that calling :meth:`env.reset()\n            <compiler_gym.envs.CompilerEnv.reset>` with one environment will\n            advance the iterator in the other. If :code:`False`, forked\n            environments will use :code:`itertools.tee()` to create a copy of\n            the iterator so that each iterator may advance independently.\n            However, this requires shared buffers between the environments which\n            can lead to memory overheads if :meth:`env.reset()\n            <compiler_gym.envs.CompilerEnv.reset>` is called many times more in\n            one environment than the other.\n        \"\"\"\n        super().__init__(env)\n        self.benchmarks = iter(benchmarks)\n        self.fork_shares_iterator = fork_shares_iterator\n\n    def reset(self, benchmark: Optional[BenchmarkLike] = None, **kwargs):\n        if benchmark is not None:\n            raise TypeError(\"Benchmark passed to IterateOverBenchmarks.reset()\")\n        benchmark: BenchmarkLike = next(self.benchmarks)\n        return self.env.reset(benchmark=benchmark)\n\n    def fork(self) -> \"IterateOverBenchmarks\":\n        if self.fork_shares_iterator:\n            other_benchmarks_iterator = self.benchmarks\n        else:\n            self.benchmarks, other_benchmarks_iterator = thread_safe_tee(\n                self.benchmarks\n            )\n        return IterateOverBenchmarks(\n            env=self.env.fork(),\n            benchmarks=other_benchmarks_iterator,\n            fork_shares_iterator=self.fork_shares_iterator,\n        )\n\n\nclass CycleOverBenchmarks(IterateOverBenchmarks):\n    \"\"\"Cycle through a list of benchmarks on each call to :meth:`reset()\n    <compiler_gym.envs.CompilerEnv.reset>`. Same as\n    :class:`IterateOverBenchmarks` except the list of benchmarks repeats once\n    exhausted.\n    \"\"\"\n\n    def __init__(\n        self,\n        env: CompilerEnv,\n        benchmarks: Iterable[BenchmarkLike],\n        fork_shares_iterator: bool = False,\n    ):\n        \"\"\"Constructor.\n\n        :param env: The environment to wrap.\n\n        :param benchmarks: An iterable sequence of benchmarks.\n\n        :param fork_shares_iterator: If :code:`True`, the :code:`benchmarks`\n            iterator will be shared by a forked environment created by\n            :meth:`env.fork() <compiler_gym.envs.CompilerEnv.fork>`. This means\n            that calling :meth:`env.reset()\n            <compiler_gym.envs.CompilerEnv.reset>` with one environment will\n            advance the iterator in the other. If :code:`False`, forked\n            environments will use :code:`itertools.tee()` to create a copy of\n            the iterator so that each iterator may advance independently.\n            However, this requires shared buffers between the environments which\n            can lead to memory overheads if :meth:`env.reset()\n            <compiler_gym.envs.CompilerEnv.reset>` is called many times more in\n            one environment than the other.\n        \"\"\"\n        super().__init__(\n            env, benchmarks=cycle(benchmarks), fork_shares_iterator=fork_shares_iterator\n        )\n\n\nclass CycleOverBenchmarksIterator(CompilerEnvWrapper):\n    \"\"\"Same as :class:`CycleOverBenchmarks\n    <compiler_gym.wrappers.CycleOverBenchmarks>` except that the user generates\n    the iterator.\n    \"\"\"\n\n    def __init__(\n        self,\n        env: CompilerEnv,\n        make_benchmark_iterator: Callable[[], Iterable[BenchmarkLike]],\n    ):\n        \"\"\"Constructor.\n\n        :param env: The environment to wrap.\n\n        :param make_benchmark_iterator: A callback that returns an iterator over\n            a sequence of benchmarks. Once the iterator is exhausted, this\n            callback is called to produce a new iterator.\n        \"\"\"\n        super().__init__(env)\n        self.make_benchmark_iterator = make_benchmark_iterator\n        self.benchmarks = iter(self.make_benchmark_iterator())\n\n    def reset(self, benchmark: Optional[BenchmarkLike] = None, **kwargs):\n        if benchmark is not None:\n            raise TypeError(\"Benchmark passed toIterateOverBenchmarks.reset()\")\n        try:\n            benchmark: BenchmarkLike = next(self.benchmarks)\n        except StopIteration:\n            self.benchmarks = iter(self.make_benchmark_iterator())\n            benchmark: BenchmarkLike = next(self.benchmarks)\n\n        return self.env.reset(benchmark=benchmark)\n\n    def fork(self) -> \"CycleOverBenchmarksIterator\":\n        return CycleOverBenchmarksIterator(\n            env=self.env.fork(),\n            make_benchmark_iterator=self.make_benchmark_iterator,\n        )\n\n\nclass RandomOrderBenchmarks(IterateOverBenchmarks):\n    \"\"\"Select randomly from a list of benchmarks on each call to :meth:`reset()\n    <compiler_gym.envs.CompilerEnv.reset>`.\n\n    .. note::\n\n        Uniform random selection is provided by evaluating the input benchmarks\n        iterator into a list and sampling randomly from the list. For very large\n        and infinite iterables of benchmarks you must use the\n        :class:`IterateOverBenchmarks\n        <compiler_gym.wrappers.IterateOverBenchmarks>` wrapper with your own\n        random sampling iterator.\n    \"\"\"\n\n    def __init__(\n        self,\n        env: CompilerEnv,\n        benchmarks: Iterable[BenchmarkLike],\n        rng: Optional[np.random.Generator] = None,\n    ):\n        \"\"\"Constructor.\n\n        :param env: The environment to wrap.\n\n        :param benchmarks: An iterable sequence of benchmarks. The entirety of\n            this input iterator is evaluated during construction.\n\n        :param rng: A random number generator to use for random benchmark\n            selection.\n        \"\"\"\n        self._all_benchmarks = list(benchmarks)\n        rng = rng or np.random.default_rng()\n        super().__init__(\n            env,\n            benchmarks=(rng.choice(self._all_benchmarks) for _ in iter(int, 1)),\n            fork_shares_iterator=True,\n        )\n\n    def fork(self) -> \"IterateOverBenchmarks\":\n        \"\"\"Fork the random order benchmark wrapper.\n\n        Note that RNG state is not copied to forked environments.\n        \"\"\"\n        return IterateOverBenchmarks(\n            env=self.env.fork(), benchmarks=self._all_benchmarks\n        )\n", "path": "compiler_gym/wrappers/datasets.py"}]}