We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Possible error in how to generate unbiased benchmarks
## 📚 Documentation
(Possible error, possibly me doing it wrong)
https://facebookresearch.github.io/CompilerGym/compiler_gym/datasets.html

The code for `random_benchmark` says:

```python
rng = np.random.default_rng()
finite_datasets = [d for d in env.datasets if len(d) != math.inf]
dataset = rng.choice(
    finite_datasets,
    p=[len(d) for d in finite_datasets]
)
dataset.random_benchmark(random_state=rng)
```

But, I think that if you call `len` on a generator, it will fail with error:
`TypeError: 'float' object cannot be interpreted as an integer`

Suggest changing it to:
`d.size != math.inf`


--- 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/datasets/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 collections import deque
6 from typing import Dict, Iterable, Optional, Set, TypeVar
7 
8 import numpy as np
9 
10 from compiler_gym.datasets.benchmark import Benchmark
11 from compiler_gym.datasets.dataset import Dataset
12 from compiler_gym.datasets.uri import BENCHMARK_URI_RE, resolve_uri_protocol
13 
14 T = TypeVar("T")
15 
16 
17 def round_robin_iterables(iters: Iterable[Iterable[T]]) -> Iterable[T]:
18     """Yield from the given iterators in round robin order."""
19     # Use a queue of iterators to iterate over. Repeatedly pop an iterator from
20     # the queue, yield the next value from it, then put it at the back of the
21     # queue. The iterator is discarded once exhausted.
22     iters = deque(iters)
23     while len(iters) > 1:
24         it = iters.popleft()
25         try:
26             yield next(it)
27             iters.append(it)
28         except StopIteration:
29             pass
30     # Once we have only a single iterator left, return it directly rather
31     # continuing with the round robin.
32     if len(iters) == 1:
33         yield from iters.popleft()
34 
35 
36 class Datasets:
37     """A collection of datasets.
38 
39     This class provides a dictionary-like interface for indexing and iterating
40     over multiple :class:`Dataset <compiler_gym.datasets.Dataset>` objects.
41     Select a dataset by URI using:
42 
43         >>> env.datasets["benchmark://cbench-v1"]
44 
45     Check whether a dataset exists using:
46 
47         >>> "benchmark://cbench-v1" in env.datasets
48         True
49 
50     Or iterate over the datasets using:
51 
52         >>> for dataset in env.datasets:
53         ...     print(dataset.name)
54         benchmark://cbench-v1
55         benchmark://github-v0
56         benchmark://npb-v0
57 
58     To select a benchmark from the datasets, use :meth:`benchmark()`:
59 
60         >>> env.datasets.benchmark("benchmark://a-v0/a")
61 
62     Use the :meth:`benchmarks()` method to iterate over every benchmark in the
63     datasets in a stable round robin order:
64 
65         >>> for benchmark in env.datasets.benchmarks():
66         ...     print(benchmark)
67         benchmark://cbench-v1/1
68         benchmark://github-v0/1
69         benchmark://npb-v0/1
70         benchmark://cbench-v1/2
71         ...
72 
73     If you want to exclude a dataset, delete it:
74 
75         >>> del env.datasets["benchmark://b-v0"]
76     """
77 
78     def __init__(
79         self,
80         datasets: Iterable[Dataset],
81     ):
82         self._datasets: Dict[str, Dataset] = {d.name: d for d in datasets}
83         self._visible_datasets: Set[str] = set(
84             name for name, dataset in self._datasets.items() if not dataset.deprecated
85         )
86 
87     def datasets(self, with_deprecated: bool = False) -> Iterable[Dataset]:
88         """Enumerate the datasets.
89 
90         Dataset order is consistent across runs.
91 
92         :param with_deprecated: If :code:`True`, include datasets that have been
93             marked as deprecated.
94 
95         :return: An iterable sequence of :meth:`Dataset
96             <compiler_gym.datasets.Dataset>` instances.
97         """
98         datasets = self._datasets.values()
99         if not with_deprecated:
100             datasets = (d for d in datasets if not d.deprecated)
101         yield from sorted(datasets, key=lambda d: (d.sort_order, d.name))
102 
103     def __iter__(self) -> Iterable[Dataset]:
104         """Iterate over the datasets.
105 
106         Dataset order is consistent across runs.
107 
108         Equivalent to :meth:`datasets.datasets()
109         <compiler_gym.datasets.Dataset.datasets>`, but without the ability to
110         iterate over the deprecated datasets.
111 
112         If the number of benchmarks in any of the datasets is infinite
113         (:code:`len(dataset) == math.inf`), the iterable returned by this method
114         will continue indefinitely.
115 
116         :return: An iterable sequence of :meth:`Dataset
117             <compiler_gym.datasets.Dataset>` instances.
118         """
119         return self.datasets()
120 
121     def dataset(self, dataset: str) -> Dataset:
122         """Get a dataset.
123 
124         Return the corresponding :meth:`Dataset
125         <compiler_gym.datasets.Dataset>`. Name lookup will succeed whether or
126         not the dataset is deprecated.
127 
128         :param dataset: A dataset name.
129 
130         :return: A :meth:`Dataset <compiler_gym.datasets.Dataset>` instance.
131 
132         :raises LookupError: If :code:`dataset` is not found.
133         """
134         dataset_name = resolve_uri_protocol(dataset)
135 
136         if dataset_name not in self._datasets:
137             raise LookupError(f"Dataset not found: {dataset_name}")
138 
139         return self._datasets[dataset_name]
140 
141     def __getitem__(self, dataset: str) -> Dataset:
142         """Lookup a dataset.
143 
144         :param dataset: A dataset name.
145 
146         :return: A :meth:`Dataset <compiler_gym.datasets.Dataset>` instance.
147 
148         :raises LookupError: If :code:`dataset` is not found.
149         """
150         return self.dataset(dataset)
151 
152     def __setitem__(self, key: str, dataset: Dataset):
153         """Add a dataset to the collection.
154 
155         :param key: The name of the dataset.
156         :param dataset: The dataset to add.
157         """
158         dataset_name = resolve_uri_protocol(key)
159 
160         self._datasets[dataset_name] = dataset
161         if not dataset.deprecated:
162             self._visible_datasets.add(dataset_name)
163 
164     def __delitem__(self, dataset: str):
165         """Remove a dataset from the collection.
166 
167         This does not affect any underlying storage used by dataset. See
168         :meth:`uninstall() <compiler_gym.datasets.Datasets.uninstall>` to clean
169         up.
170 
171         :param dataset: The name of a dataset.
172 
173         :return: :code:`True` if the dataset was removed, :code:`False` if it
174             was already removed.
175         """
176         dataset_name = resolve_uri_protocol(dataset)
177         if dataset_name in self._visible_datasets:
178             self._visible_datasets.remove(dataset_name)
179         del self._datasets[dataset_name]
180 
181     def __contains__(self, dataset: str) -> bool:
182         """Returns whether the dataset is contained."""
183         try:
184             self.dataset(dataset)
185             return True
186         except LookupError:
187             return False
188 
189     def benchmarks(self, with_deprecated: bool = False) -> Iterable[Benchmark]:
190         """Enumerate the (possibly infinite) benchmarks lazily.
191 
192         Benchmarks order is consistent across runs. One benchmark from each
193         dataset is returned in round robin order until all datasets have been
194         fully enumerated. The order of :meth:`benchmarks()
195         <compiler_gym.datasets.Datasets.benchmarks>` and :meth:`benchmark_uris()
196         <compiler_gym.datasets.Datasets.benchmark_uris>` is the same.
197 
198         If the number of benchmarks in any of the datasets is infinite
199         (:code:`len(dataset) == math.inf`), the iterable returned by this method
200         will continue indefinitely.
201 
202         :param with_deprecated: If :code:`True`, include benchmarks from
203             datasets that have been marked deprecated.
204 
205         :return: An iterable sequence of :class:`Benchmark
206             <compiler_gym.datasets.Benchmark>` instances.
207         """
208         return round_robin_iterables(
209             (d.benchmarks() for d in self.datasets(with_deprecated=with_deprecated))
210         )
211 
212     def benchmark_uris(self, with_deprecated: bool = False) -> Iterable[str]:
213         """Enumerate the (possibly infinite) benchmark URIs.
214 
215         Benchmark URI order is consistent across runs. URIs from datasets are
216         returned in round robin order. The order of :meth:`benchmarks()
217         <compiler_gym.datasets.Datasets.benchmarks>` and :meth:`benchmark_uris()
218         <compiler_gym.datasets.Datasets.benchmark_uris>` is the same.
219 
220         If the number of benchmarks in any of the datasets is infinite
221         (:code:`len(dataset) == math.inf`), the iterable returned by this method
222         will continue indefinitely.
223 
224         :param with_deprecated: If :code:`True`, include benchmarks from
225             datasets that have been marked deprecated.
226 
227         :return: An iterable sequence of benchmark URI strings.
228         """
229         return round_robin_iterables(
230             (d.benchmark_uris() for d in self.datasets(with_deprecated=with_deprecated))
231         )
232 
233     def benchmark(self, uri: str) -> Benchmark:
234         """Select a benchmark.
235 
236         Returns the corresponding :class:`Benchmark
237         <compiler_gym.datasets.Benchmark>`, regardless of whether the containing
238         dataset is installed or deprecated.
239 
240         :param uri: The URI of the benchmark to return.
241 
242         :return: A :class:`Benchmark <compiler_gym.datasets.Benchmark>`
243             instance.
244         """
245         uri = resolve_uri_protocol(uri)
246 
247         match = BENCHMARK_URI_RE.match(uri)
248         if not match:
249             raise ValueError(f"Invalid benchmark URI: '{uri}'")
250 
251         dataset_name = match.group("dataset")
252         dataset = self._datasets[dataset_name]
253 
254         return dataset.benchmark(uri)
255 
256     def random_benchmark(
257         self, random_state: Optional[np.random.Generator] = None
258     ) -> Benchmark:
259         """Select a benchmark randomly.
260 
261         First, a dataset is selected uniformly randomly using
262         :code:`random_state.choice(list(datasets))`. The
263         :meth:`random_benchmark()
264         <compiler_gym.datasets.Dataset.random_benchmark>` method of that dataset
265         is then called to select a benchmark.
266 
267         Note that the distribution of benchmarks selected by this method is not
268         biased by the size of each dataset, since datasets are selected
269         uniformly. This means that datasets with a small number of benchmarks
270         will be overrepresented compared to datasets with many benchmarks. To
271         correct for this bias, use the number of benchmarks in each dataset as
272         a weight for the random selection:
273 
274             >>> rng = np.random.default_rng()
275             >>> finite_datasets = [d for d in env.datasets if len(d) != math.inf]
276             >>> dataset = rng.choice(
277                 finite_datasets,
278                 p=[len(d) for d in finite_datasets]
279             )
280             >>> dataset.random_benchmark(random_state=rng)
281 
282         :param random_state: A random number generator. If not provided, a
283             default :code:`np.random.default_rng()` is used.
284 
285         :return: A :class:`Benchmark <compiler_gym.datasets.Benchmark>`
286             instance.
287         """
288         random_state = random_state or np.random.default_rng()
289         dataset = random_state.choice(list(self._visible_datasets))
290         return self[dataset].random_benchmark(random_state=random_state)
291 
292     @property
293     def size(self) -> int:
294         return len(self._visible_datasets)
295 
296     def __len__(self) -> int:
297         """The number of datasets in the collection."""
298         return self.size
299
```
--- 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/datasets/datasets.py b/compiler_gym/datasets/datasets.py
--- a/compiler_gym/datasets/datasets.py
+++ b/compiler_gym/datasets/datasets.py
@@ -3,7 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 from collections import deque
-from typing import Dict, Iterable, Optional, Set, TypeVar
+from typing import Dict, Iterable, List, Optional, Set, TypeVar
 
 import numpy as np
 
@@ -254,39 +254,64 @@
         return dataset.benchmark(uri)
 
     def random_benchmark(
-        self, random_state: Optional[np.random.Generator] = None
+        self,
+        random_state: Optional[np.random.Generator] = None,
+        weighted: bool = False,
+        weights: Optional[Dict[str, float]] = None,
     ) -> Benchmark:
         """Select a benchmark randomly.
 
-        First, a dataset is selected uniformly randomly using
-        :code:`random_state.choice(list(datasets))`. The
+        First, a dataset is selected randomly using
+        :code:`random_state.choice(list(datasets))`. Then the
         :meth:`random_benchmark()
-        <compiler_gym.datasets.Dataset.random_benchmark>` method of that dataset
-        is then called to select a benchmark.
-
-        Note that the distribution of benchmarks selected by this method is not
-        biased by the size of each dataset, since datasets are selected
-        uniformly. This means that datasets with a small number of benchmarks
-        will be overrepresented compared to datasets with many benchmarks. To
-        correct for this bias, use the number of benchmarks in each dataset as
-        a weight for the random selection:
-
-            >>> rng = np.random.default_rng()
-            >>> finite_datasets = [d for d in env.datasets if len(d) != math.inf]
-            >>> dataset = rng.choice(
-                finite_datasets,
-                p=[len(d) for d in finite_datasets]
-            )
-            >>> dataset.random_benchmark(random_state=rng)
+        <compiler_gym.datasets.Dataset.random_benchmark>` method of the chosen
+        dataset is called to select a benchmark.
+
+        By default datasets are selected uniformly randomly. This means that
+        datasets with a small number of benchmarks will be overrepresented
+        compared to datasets with many benchmarks. To correct for this bias pass
+        the argument :code:`weighted=True`, which weights the dataset choice by
+        the number of benchmarks in each dataset, equivalent to:
+
+            >>> random.choices(datasets, weights=[len(p) for p in datasets])
+
+        Weighting the choice of datasets by their size means that datasets with
+        infinite sizes (such as random program generators) will be excluded from
+        sampling as their size is :code:`0`. To override the weights of datasets
+        pass a :code:`weights` mapping:
+
+            >>> env.datasets.random_benchmark(weighted=True, weights={
+                "benchmark://dataset-v0": 10,
+                "benchmark://another-dataset-v0": 555,
+            })
 
         :param random_state: A random number generator. If not provided, a
             default :code:`np.random.default_rng()` is used.
 
+        :param weighted: If set, weight the choice of dataset by the number of
+            benchmarks in each dataset, or the value specified in the
+            :code:`weights` mapping.
+
+        :param weights: An optional mapping from dataset URI to the weight to
+            use when :code:`weighted=True`. This overrides the default value of
+            using the dataset size.
+
         :return: A :class:`Benchmark <compiler_gym.datasets.Benchmark>`
             instance.
         """
         random_state = random_state or np.random.default_rng()
-        dataset = random_state.choice(list(self._visible_datasets))
+        datasets: List[str] = list(self._visible_datasets)
+        # Assume weighted=True if weights dictionary is specified.
+        weighted = weighted or weights
+
+        if weighted:
+            weights: Dict[str, float] = weights or {}
+            w: List[float] = np.array(
+                [weights.get(d, self[d].size) for d in datasets], dtype=float
+            )
+            dataset = random_state.choice(datasets, p=w / w.sum())
+        else:
+            dataset = random_state.choice(datasets)
         return self[dataset].random_benchmark(random_state=random_state)
 
     @property

{"golden_diff": "diff --git a/compiler_gym/datasets/datasets.py b/compiler_gym/datasets/datasets.py\n--- a/compiler_gym/datasets/datasets.py\n+++ b/compiler_gym/datasets/datasets.py\n@@ -3,7 +3,7 @@\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 from collections import deque\n-from typing import Dict, Iterable, Optional, Set, TypeVar\n+from typing import Dict, Iterable, List, Optional, Set, TypeVar\n \n import numpy as np\n \n@@ -254,39 +254,64 @@\n         return dataset.benchmark(uri)\n \n     def random_benchmark(\n-        self, random_state: Optional[np.random.Generator] = None\n+        self,\n+        random_state: Optional[np.random.Generator] = None,\n+        weighted: bool = False,\n+        weights: Optional[Dict[str, float]] = None,\n     ) -> Benchmark:\n         \"\"\"Select a benchmark randomly.\n \n-        First, a dataset is selected uniformly randomly using\n-        :code:`random_state.choice(list(datasets))`. The\n+        First, a dataset is selected randomly using\n+        :code:`random_state.choice(list(datasets))`. Then the\n         :meth:`random_benchmark()\n-        <compiler_gym.datasets.Dataset.random_benchmark>` method of that dataset\n-        is then called to select a benchmark.\n-\n-        Note that the distribution of benchmarks selected by this method is not\n-        biased by the size of each dataset, since datasets are selected\n-        uniformly. This means that datasets with a small number of benchmarks\n-        will be overrepresented compared to datasets with many benchmarks. To\n-        correct for this bias, use the number of benchmarks in each dataset as\n-        a weight for the random selection:\n-\n-            >>> rng = np.random.default_rng()\n-            >>> finite_datasets = [d for d in env.datasets if len(d) != math.inf]\n-            >>> dataset = rng.choice(\n-                finite_datasets,\n-                p=[len(d) for d in finite_datasets]\n-            )\n-            >>> dataset.random_benchmark(random_state=rng)\n+        <compiler_gym.datasets.Dataset.random_benchmark>` method of the chosen\n+        dataset is called to select a benchmark.\n+\n+        By default datasets are selected uniformly randomly. This means that\n+        datasets with a small number of benchmarks will be overrepresented\n+        compared to datasets with many benchmarks. To correct for this bias pass\n+        the argument :code:`weighted=True`, which weights the dataset choice by\n+        the number of benchmarks in each dataset, equivalent to:\n+\n+            >>> random.choices(datasets, weights=[len(p) for p in datasets])\n+\n+        Weighting the choice of datasets by their size means that datasets with\n+        infinite sizes (such as random program generators) will be excluded from\n+        sampling as their size is :code:`0`. To override the weights of datasets\n+        pass a :code:`weights` mapping:\n+\n+            >>> env.datasets.random_benchmark(weighted=True, weights={\n+                \"benchmark://dataset-v0\": 10,\n+                \"benchmark://another-dataset-v0\": 555,\n+            })\n \n         :param random_state: A random number generator. If not provided, a\n             default :code:`np.random.default_rng()` is used.\n \n+        :param weighted: If set, weight the choice of dataset by the number of\n+            benchmarks in each dataset, or the value specified in the\n+            :code:`weights` mapping.\n+\n+        :param weights: An optional mapping from dataset URI to the weight to\n+            use when :code:`weighted=True`. This overrides the default value of\n+            using the dataset size.\n+\n         :return: A :class:`Benchmark <compiler_gym.datasets.Benchmark>`\n             instance.\n         \"\"\"\n         random_state = random_state or np.random.default_rng()\n-        dataset = random_state.choice(list(self._visible_datasets))\n+        datasets: List[str] = list(self._visible_datasets)\n+        # Assume weighted=True if weights dictionary is specified.\n+        weighted = weighted or weights\n+\n+        if weighted:\n+            weights: Dict[str, float] = weights or {}\n+            w: List[float] = np.array(\n+                [weights.get(d, self[d].size) for d in datasets], dtype=float\n+            )\n+            dataset = random_state.choice(datasets, p=w / w.sum())\n+        else:\n+            dataset = random_state.choice(datasets)\n         return self[dataset].random_benchmark(random_state=random_state)\n \n     @property\n", "issue": "Possible error in how to generate unbiased benchmarks\n## \ud83d\udcda Documentation\r\n(Possible error, possibly me doing it wrong)\r\nhttps://facebookresearch.github.io/CompilerGym/compiler_gym/datasets.html\r\n\r\nThe code for `random_benchmark` says:\r\n\r\n```python\r\nrng = np.random.default_rng()\r\nfinite_datasets = [d for d in env.datasets if len(d) != math.inf]\r\ndataset = rng.choice(\r\n    finite_datasets,\r\n    p=[len(d) for d in finite_datasets]\r\n)\r\ndataset.random_benchmark(random_state=rng)\r\n```\r\n\r\nBut, I think that if you call `len` on a generator, it will fail with error:\r\n`TypeError: 'float' object cannot be interpreted as an integer`\r\n\r\nSuggest changing it to:\r\n`d.size != math.inf`\r\n\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.\nfrom collections import deque\nfrom typing import Dict, Iterable, Optional, Set, TypeVar\n\nimport numpy as np\n\nfrom compiler_gym.datasets.benchmark import Benchmark\nfrom compiler_gym.datasets.dataset import Dataset\nfrom compiler_gym.datasets.uri import BENCHMARK_URI_RE, resolve_uri_protocol\n\nT = TypeVar(\"T\")\n\n\ndef round_robin_iterables(iters: Iterable[Iterable[T]]) -> Iterable[T]:\n    \"\"\"Yield from the given iterators in round robin order.\"\"\"\n    # Use a queue of iterators to iterate over. Repeatedly pop an iterator from\n    # the queue, yield the next value from it, then put it at the back of the\n    # queue. The iterator is discarded once exhausted.\n    iters = deque(iters)\n    while len(iters) > 1:\n        it = iters.popleft()\n        try:\n            yield next(it)\n            iters.append(it)\n        except StopIteration:\n            pass\n    # Once we have only a single iterator left, return it directly rather\n    # continuing with the round robin.\n    if len(iters) == 1:\n        yield from iters.popleft()\n\n\nclass Datasets:\n    \"\"\"A collection of datasets.\n\n    This class provides a dictionary-like interface for indexing and iterating\n    over multiple :class:`Dataset <compiler_gym.datasets.Dataset>` objects.\n    Select a dataset by URI using:\n\n        >>> env.datasets[\"benchmark://cbench-v1\"]\n\n    Check whether a dataset exists using:\n\n        >>> \"benchmark://cbench-v1\" in env.datasets\n        True\n\n    Or iterate over the datasets using:\n\n        >>> for dataset in env.datasets:\n        ...     print(dataset.name)\n        benchmark://cbench-v1\n        benchmark://github-v0\n        benchmark://npb-v0\n\n    To select a benchmark from the datasets, use :meth:`benchmark()`:\n\n        >>> env.datasets.benchmark(\"benchmark://a-v0/a\")\n\n    Use the :meth:`benchmarks()` method to iterate over every benchmark in the\n    datasets in a stable round robin order:\n\n        >>> for benchmark in env.datasets.benchmarks():\n        ...     print(benchmark)\n        benchmark://cbench-v1/1\n        benchmark://github-v0/1\n        benchmark://npb-v0/1\n        benchmark://cbench-v1/2\n        ...\n\n    If you want to exclude a dataset, delete it:\n\n        >>> del env.datasets[\"benchmark://b-v0\"]\n    \"\"\"\n\n    def __init__(\n        self,\n        datasets: Iterable[Dataset],\n    ):\n        self._datasets: Dict[str, Dataset] = {d.name: d for d in datasets}\n        self._visible_datasets: Set[str] = set(\n            name for name, dataset in self._datasets.items() if not dataset.deprecated\n        )\n\n    def datasets(self, with_deprecated: bool = False) -> Iterable[Dataset]:\n        \"\"\"Enumerate the datasets.\n\n        Dataset order is consistent across runs.\n\n        :param with_deprecated: If :code:`True`, include datasets that have been\n            marked as deprecated.\n\n        :return: An iterable sequence of :meth:`Dataset\n            <compiler_gym.datasets.Dataset>` instances.\n        \"\"\"\n        datasets = self._datasets.values()\n        if not with_deprecated:\n            datasets = (d for d in datasets if not d.deprecated)\n        yield from sorted(datasets, key=lambda d: (d.sort_order, d.name))\n\n    def __iter__(self) -> Iterable[Dataset]:\n        \"\"\"Iterate over the datasets.\n\n        Dataset order is consistent across runs.\n\n        Equivalent to :meth:`datasets.datasets()\n        <compiler_gym.datasets.Dataset.datasets>`, but without the ability to\n        iterate over the deprecated datasets.\n\n        If the number of benchmarks in any of the datasets is infinite\n        (:code:`len(dataset) == math.inf`), the iterable returned by this method\n        will continue indefinitely.\n\n        :return: An iterable sequence of :meth:`Dataset\n            <compiler_gym.datasets.Dataset>` instances.\n        \"\"\"\n        return self.datasets()\n\n    def dataset(self, dataset: str) -> Dataset:\n        \"\"\"Get a dataset.\n\n        Return the corresponding :meth:`Dataset\n        <compiler_gym.datasets.Dataset>`. Name lookup will succeed whether or\n        not the dataset is deprecated.\n\n        :param dataset: A dataset name.\n\n        :return: A :meth:`Dataset <compiler_gym.datasets.Dataset>` instance.\n\n        :raises LookupError: If :code:`dataset` is not found.\n        \"\"\"\n        dataset_name = resolve_uri_protocol(dataset)\n\n        if dataset_name not in self._datasets:\n            raise LookupError(f\"Dataset not found: {dataset_name}\")\n\n        return self._datasets[dataset_name]\n\n    def __getitem__(self, dataset: str) -> Dataset:\n        \"\"\"Lookup a dataset.\n\n        :param dataset: A dataset name.\n\n        :return: A :meth:`Dataset <compiler_gym.datasets.Dataset>` instance.\n\n        :raises LookupError: If :code:`dataset` is not found.\n        \"\"\"\n        return self.dataset(dataset)\n\n    def __setitem__(self, key: str, dataset: Dataset):\n        \"\"\"Add a dataset to the collection.\n\n        :param key: The name of the dataset.\n        :param dataset: The dataset to add.\n        \"\"\"\n        dataset_name = resolve_uri_protocol(key)\n\n        self._datasets[dataset_name] = dataset\n        if not dataset.deprecated:\n            self._visible_datasets.add(dataset_name)\n\n    def __delitem__(self, dataset: str):\n        \"\"\"Remove a dataset from the collection.\n\n        This does not affect any underlying storage used by dataset. See\n        :meth:`uninstall() <compiler_gym.datasets.Datasets.uninstall>` to clean\n        up.\n\n        :param dataset: The name of a dataset.\n\n        :return: :code:`True` if the dataset was removed, :code:`False` if it\n            was already removed.\n        \"\"\"\n        dataset_name = resolve_uri_protocol(dataset)\n        if dataset_name in self._visible_datasets:\n            self._visible_datasets.remove(dataset_name)\n        del self._datasets[dataset_name]\n\n    def __contains__(self, dataset: str) -> bool:\n        \"\"\"Returns whether the dataset is contained.\"\"\"\n        try:\n            self.dataset(dataset)\n            return True\n        except LookupError:\n            return False\n\n    def benchmarks(self, with_deprecated: bool = False) -> Iterable[Benchmark]:\n        \"\"\"Enumerate the (possibly infinite) benchmarks lazily.\n\n        Benchmarks order is consistent across runs. One benchmark from each\n        dataset is returned in round robin order until all datasets have been\n        fully enumerated. The order of :meth:`benchmarks()\n        <compiler_gym.datasets.Datasets.benchmarks>` and :meth:`benchmark_uris()\n        <compiler_gym.datasets.Datasets.benchmark_uris>` is the same.\n\n        If the number of benchmarks in any of the datasets is infinite\n        (:code:`len(dataset) == math.inf`), the iterable returned by this method\n        will continue indefinitely.\n\n        :param with_deprecated: If :code:`True`, include benchmarks from\n            datasets that have been marked deprecated.\n\n        :return: An iterable sequence of :class:`Benchmark\n            <compiler_gym.datasets.Benchmark>` instances.\n        \"\"\"\n        return round_robin_iterables(\n            (d.benchmarks() for d in self.datasets(with_deprecated=with_deprecated))\n        )\n\n    def benchmark_uris(self, with_deprecated: bool = False) -> Iterable[str]:\n        \"\"\"Enumerate the (possibly infinite) benchmark URIs.\n\n        Benchmark URI order is consistent across runs. URIs from datasets are\n        returned in round robin order. The order of :meth:`benchmarks()\n        <compiler_gym.datasets.Datasets.benchmarks>` and :meth:`benchmark_uris()\n        <compiler_gym.datasets.Datasets.benchmark_uris>` is the same.\n\n        If the number of benchmarks in any of the datasets is infinite\n        (:code:`len(dataset) == math.inf`), the iterable returned by this method\n        will continue indefinitely.\n\n        :param with_deprecated: If :code:`True`, include benchmarks from\n            datasets that have been marked deprecated.\n\n        :return: An iterable sequence of benchmark URI strings.\n        \"\"\"\n        return round_robin_iterables(\n            (d.benchmark_uris() for d in self.datasets(with_deprecated=with_deprecated))\n        )\n\n    def benchmark(self, uri: str) -> Benchmark:\n        \"\"\"Select a benchmark.\n\n        Returns the corresponding :class:`Benchmark\n        <compiler_gym.datasets.Benchmark>`, regardless of whether the containing\n        dataset is installed or deprecated.\n\n        :param uri: The URI of the benchmark to return.\n\n        :return: A :class:`Benchmark <compiler_gym.datasets.Benchmark>`\n            instance.\n        \"\"\"\n        uri = resolve_uri_protocol(uri)\n\n        match = BENCHMARK_URI_RE.match(uri)\n        if not match:\n            raise ValueError(f\"Invalid benchmark URI: '{uri}'\")\n\n        dataset_name = match.group(\"dataset\")\n        dataset = self._datasets[dataset_name]\n\n        return dataset.benchmark(uri)\n\n    def random_benchmark(\n        self, random_state: Optional[np.random.Generator] = None\n    ) -> Benchmark:\n        \"\"\"Select a benchmark randomly.\n\n        First, a dataset is selected uniformly randomly using\n        :code:`random_state.choice(list(datasets))`. The\n        :meth:`random_benchmark()\n        <compiler_gym.datasets.Dataset.random_benchmark>` method of that dataset\n        is then called to select a benchmark.\n\n        Note that the distribution of benchmarks selected by this method is not\n        biased by the size of each dataset, since datasets are selected\n        uniformly. This means that datasets with a small number of benchmarks\n        will be overrepresented compared to datasets with many benchmarks. To\n        correct for this bias, use the number of benchmarks in each dataset as\n        a weight for the random selection:\n\n            >>> rng = np.random.default_rng()\n            >>> finite_datasets = [d for d in env.datasets if len(d) != math.inf]\n            >>> dataset = rng.choice(\n                finite_datasets,\n                p=[len(d) for d in finite_datasets]\n            )\n            >>> dataset.random_benchmark(random_state=rng)\n\n        :param random_state: A random number generator. If not provided, a\n            default :code:`np.random.default_rng()` is used.\n\n        :return: A :class:`Benchmark <compiler_gym.datasets.Benchmark>`\n            instance.\n        \"\"\"\n        random_state = random_state or np.random.default_rng()\n        dataset = random_state.choice(list(self._visible_datasets))\n        return self[dataset].random_benchmark(random_state=random_state)\n\n    @property\n    def size(self) -> int:\n        return len(self._visible_datasets)\n\n    def __len__(self) -> int:\n        \"\"\"The number of datasets in the collection.\"\"\"\n        return self.size\n", "path": "compiler_gym/datasets/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.\nfrom collections import deque\nfrom typing import Dict, Iterable, List, Optional, Set, TypeVar\n\nimport numpy as np\n\nfrom compiler_gym.datasets.benchmark import Benchmark\nfrom compiler_gym.datasets.dataset import Dataset\nfrom compiler_gym.datasets.uri import BENCHMARK_URI_RE, resolve_uri_protocol\n\nT = TypeVar(\"T\")\n\n\ndef round_robin_iterables(iters: Iterable[Iterable[T]]) -> Iterable[T]:\n    \"\"\"Yield from the given iterators in round robin order.\"\"\"\n    # Use a queue of iterators to iterate over. Repeatedly pop an iterator from\n    # the queue, yield the next value from it, then put it at the back of the\n    # queue. The iterator is discarded once exhausted.\n    iters = deque(iters)\n    while len(iters) > 1:\n        it = iters.popleft()\n        try:\n            yield next(it)\n            iters.append(it)\n        except StopIteration:\n            pass\n    # Once we have only a single iterator left, return it directly rather\n    # continuing with the round robin.\n    if len(iters) == 1:\n        yield from iters.popleft()\n\n\nclass Datasets:\n    \"\"\"A collection of datasets.\n\n    This class provides a dictionary-like interface for indexing and iterating\n    over multiple :class:`Dataset <compiler_gym.datasets.Dataset>` objects.\n    Select a dataset by URI using:\n\n        >>> env.datasets[\"benchmark://cbench-v1\"]\n\n    Check whether a dataset exists using:\n\n        >>> \"benchmark://cbench-v1\" in env.datasets\n        True\n\n    Or iterate over the datasets using:\n\n        >>> for dataset in env.datasets:\n        ...     print(dataset.name)\n        benchmark://cbench-v1\n        benchmark://github-v0\n        benchmark://npb-v0\n\n    To select a benchmark from the datasets, use :meth:`benchmark()`:\n\n        >>> env.datasets.benchmark(\"benchmark://a-v0/a\")\n\n    Use the :meth:`benchmarks()` method to iterate over every benchmark in the\n    datasets in a stable round robin order:\n\n        >>> for benchmark in env.datasets.benchmarks():\n        ...     print(benchmark)\n        benchmark://cbench-v1/1\n        benchmark://github-v0/1\n        benchmark://npb-v0/1\n        benchmark://cbench-v1/2\n        ...\n\n    If you want to exclude a dataset, delete it:\n\n        >>> del env.datasets[\"benchmark://b-v0\"]\n    \"\"\"\n\n    def __init__(\n        self,\n        datasets: Iterable[Dataset],\n    ):\n        self._datasets: Dict[str, Dataset] = {d.name: d for d in datasets}\n        self._visible_datasets: Set[str] = set(\n            name for name, dataset in self._datasets.items() if not dataset.deprecated\n        )\n\n    def datasets(self, with_deprecated: bool = False) -> Iterable[Dataset]:\n        \"\"\"Enumerate the datasets.\n\n        Dataset order is consistent across runs.\n\n        :param with_deprecated: If :code:`True`, include datasets that have been\n            marked as deprecated.\n\n        :return: An iterable sequence of :meth:`Dataset\n            <compiler_gym.datasets.Dataset>` instances.\n        \"\"\"\n        datasets = self._datasets.values()\n        if not with_deprecated:\n            datasets = (d for d in datasets if not d.deprecated)\n        yield from sorted(datasets, key=lambda d: (d.sort_order, d.name))\n\n    def __iter__(self) -> Iterable[Dataset]:\n        \"\"\"Iterate over the datasets.\n\n        Dataset order is consistent across runs.\n\n        Equivalent to :meth:`datasets.datasets()\n        <compiler_gym.datasets.Dataset.datasets>`, but without the ability to\n        iterate over the deprecated datasets.\n\n        If the number of benchmarks in any of the datasets is infinite\n        (:code:`len(dataset) == math.inf`), the iterable returned by this method\n        will continue indefinitely.\n\n        :return: An iterable sequence of :meth:`Dataset\n            <compiler_gym.datasets.Dataset>` instances.\n        \"\"\"\n        return self.datasets()\n\n    def dataset(self, dataset: str) -> Dataset:\n        \"\"\"Get a dataset.\n\n        Return the corresponding :meth:`Dataset\n        <compiler_gym.datasets.Dataset>`. Name lookup will succeed whether or\n        not the dataset is deprecated.\n\n        :param dataset: A dataset name.\n\n        :return: A :meth:`Dataset <compiler_gym.datasets.Dataset>` instance.\n\n        :raises LookupError: If :code:`dataset` is not found.\n        \"\"\"\n        dataset_name = resolve_uri_protocol(dataset)\n\n        if dataset_name not in self._datasets:\n            raise LookupError(f\"Dataset not found: {dataset_name}\")\n\n        return self._datasets[dataset_name]\n\n    def __getitem__(self, dataset: str) -> Dataset:\n        \"\"\"Lookup a dataset.\n\n        :param dataset: A dataset name.\n\n        :return: A :meth:`Dataset <compiler_gym.datasets.Dataset>` instance.\n\n        :raises LookupError: If :code:`dataset` is not found.\n        \"\"\"\n        return self.dataset(dataset)\n\n    def __setitem__(self, key: str, dataset: Dataset):\n        \"\"\"Add a dataset to the collection.\n\n        :param key: The name of the dataset.\n        :param dataset: The dataset to add.\n        \"\"\"\n        dataset_name = resolve_uri_protocol(key)\n\n        self._datasets[dataset_name] = dataset\n        if not dataset.deprecated:\n            self._visible_datasets.add(dataset_name)\n\n    def __delitem__(self, dataset: str):\n        \"\"\"Remove a dataset from the collection.\n\n        This does not affect any underlying storage used by dataset. See\n        :meth:`uninstall() <compiler_gym.datasets.Datasets.uninstall>` to clean\n        up.\n\n        :param dataset: The name of a dataset.\n\n        :return: :code:`True` if the dataset was removed, :code:`False` if it\n            was already removed.\n        \"\"\"\n        dataset_name = resolve_uri_protocol(dataset)\n        if dataset_name in self._visible_datasets:\n            self._visible_datasets.remove(dataset_name)\n        del self._datasets[dataset_name]\n\n    def __contains__(self, dataset: str) -> bool:\n        \"\"\"Returns whether the dataset is contained.\"\"\"\n        try:\n            self.dataset(dataset)\n            return True\n        except LookupError:\n            return False\n\n    def benchmarks(self, with_deprecated: bool = False) -> Iterable[Benchmark]:\n        \"\"\"Enumerate the (possibly infinite) benchmarks lazily.\n\n        Benchmarks order is consistent across runs. One benchmark from each\n        dataset is returned in round robin order until all datasets have been\n        fully enumerated. The order of :meth:`benchmarks()\n        <compiler_gym.datasets.Datasets.benchmarks>` and :meth:`benchmark_uris()\n        <compiler_gym.datasets.Datasets.benchmark_uris>` is the same.\n\n        If the number of benchmarks in any of the datasets is infinite\n        (:code:`len(dataset) == math.inf`), the iterable returned by this method\n        will continue indefinitely.\n\n        :param with_deprecated: If :code:`True`, include benchmarks from\n            datasets that have been marked deprecated.\n\n        :return: An iterable sequence of :class:`Benchmark\n            <compiler_gym.datasets.Benchmark>` instances.\n        \"\"\"\n        return round_robin_iterables(\n            (d.benchmarks() for d in self.datasets(with_deprecated=with_deprecated))\n        )\n\n    def benchmark_uris(self, with_deprecated: bool = False) -> Iterable[str]:\n        \"\"\"Enumerate the (possibly infinite) benchmark URIs.\n\n        Benchmark URI order is consistent across runs. URIs from datasets are\n        returned in round robin order. The order of :meth:`benchmarks()\n        <compiler_gym.datasets.Datasets.benchmarks>` and :meth:`benchmark_uris()\n        <compiler_gym.datasets.Datasets.benchmark_uris>` is the same.\n\n        If the number of benchmarks in any of the datasets is infinite\n        (:code:`len(dataset) == math.inf`), the iterable returned by this method\n        will continue indefinitely.\n\n        :param with_deprecated: If :code:`True`, include benchmarks from\n            datasets that have been marked deprecated.\n\n        :return: An iterable sequence of benchmark URI strings.\n        \"\"\"\n        return round_robin_iterables(\n            (d.benchmark_uris() for d in self.datasets(with_deprecated=with_deprecated))\n        )\n\n    def benchmark(self, uri: str) -> Benchmark:\n        \"\"\"Select a benchmark.\n\n        Returns the corresponding :class:`Benchmark\n        <compiler_gym.datasets.Benchmark>`, regardless of whether the containing\n        dataset is installed or deprecated.\n\n        :param uri: The URI of the benchmark to return.\n\n        :return: A :class:`Benchmark <compiler_gym.datasets.Benchmark>`\n            instance.\n        \"\"\"\n        uri = resolve_uri_protocol(uri)\n\n        match = BENCHMARK_URI_RE.match(uri)\n        if not match:\n            raise ValueError(f\"Invalid benchmark URI: '{uri}'\")\n\n        dataset_name = match.group(\"dataset\")\n        dataset = self._datasets[dataset_name]\n\n        return dataset.benchmark(uri)\n\n    def random_benchmark(\n        self,\n        random_state: Optional[np.random.Generator] = None,\n        weighted: bool = False,\n        weights: Optional[Dict[str, float]] = None,\n    ) -> Benchmark:\n        \"\"\"Select a benchmark randomly.\n\n        First, a dataset is selected randomly using\n        :code:`random_state.choice(list(datasets))`. Then the\n        :meth:`random_benchmark()\n        <compiler_gym.datasets.Dataset.random_benchmark>` method of the chosen\n        dataset is called to select a benchmark.\n\n        By default datasets are selected uniformly randomly. This means that\n        datasets with a small number of benchmarks will be overrepresented\n        compared to datasets with many benchmarks. To correct for this bias pass\n        the argument :code:`weighted=True`, which weights the dataset choice by\n        the number of benchmarks in each dataset, equivalent to:\n\n            >>> random.choices(datasets, weights=[len(p) for p in datasets])\n\n        Weighting the choice of datasets by their size means that datasets with\n        infinite sizes (such as random program generators) will be excluded from\n        sampling as their size is :code:`0`. To override the weights of datasets\n        pass a :code:`weights` mapping:\n\n            >>> env.datasets.random_benchmark(weighted=True, weights={\n                \"benchmark://dataset-v0\": 10,\n                \"benchmark://another-dataset-v0\": 555,\n            })\n\n        :param random_state: A random number generator. If not provided, a\n            default :code:`np.random.default_rng()` is used.\n\n        :param weighted: If set, weight the choice of dataset by the number of\n            benchmarks in each dataset, or the value specified in the\n            :code:`weights` mapping.\n\n        :param weights: An optional mapping from dataset URI to the weight to\n            use when :code:`weighted=True`. This overrides the default value of\n            using the dataset size.\n\n        :return: A :class:`Benchmark <compiler_gym.datasets.Benchmark>`\n            instance.\n        \"\"\"\n        random_state = random_state or np.random.default_rng()\n        datasets: List[str] = list(self._visible_datasets)\n        # Assume weighted=True if weights dictionary is specified.\n        weighted = weighted or weights\n\n        if weighted:\n            weights: Dict[str, float] = weights or {}\n            w: List[float] = np.array(\n                [weights.get(d, self[d].size) for d in datasets], dtype=float\n            )\n            dataset = random_state.choice(datasets, p=w / w.sum())\n        else:\n            dataset = random_state.choice(datasets)\n        return self[dataset].random_benchmark(random_state=random_state)\n\n    @property\n    def size(self) -> int:\n        return len(self._visible_datasets)\n\n    def __len__(self) -> int:\n        \"\"\"The number of datasets in the collection.\"\"\"\n        return self.size\n", "path": "compiler_gym/datasets/datasets.py"}]}