Add SamplerWithoutReplacement to corpus.py. (#259)

* Add SamplerWithoutReplacement to corpus.py.

* Fix function signature for type annotation.

* Yapf.

* Fix local_data_collector_test.py

* Fix construction of sampler in local_data_collector_test.py.

* Fix formatting.

* Yapf.

* Re-add multiprocessing include.

Author: jacob-hegna

compiler_opt/rl/corpus.py

@@ -20,7 +20,7 @@
 
 from absl import logging
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type
 
 import json
 import os
@@ -126,13 +126,17 @@ class Sampler(metaclass=abc.ABCMeta):
   """Corpus sampler abstraction."""
 
   @abc.abstractmethod
-  def __call__(self,
-               module_specs: Tuple[ModuleSpec],
-               k: int,
-               n: int = 20) -> List[ModuleSpec]:
+  def __init__(self, module_specs: Tuple[ModuleSpec]):
+    self._module_specs = module_specs
+
+  @abc.abstractmethod
+  def reset(self):
+    pass
+
+  @abc.abstractmethod
+  def __call__(self, k: int, n: int = 20) -> List[ModuleSpec]:
     """
     Args:
-      module_specs: list of module_specs to sample from
       k: number of modules to sample
       n: number of buckets to use
     """
@@ -144,13 +148,14 @@ class SamplerBucketRoundRobin(Sampler):
   round-robin order. The buckets are sequential sections of module_specs of
   roughly equal lengths."""
 
-  def __init__(self):
+  def __init__(self, module_specs: Tuple[ModuleSpec]):
     self._ranges = {}
+    super().__init__(module_specs)
 
-  def __call__(self,
-               module_specs: Tuple[ModuleSpec],
-               k: int,
-               n: int = 20) -> List[ModuleSpec]:
+  def reset(self):
+    pass
+
+  def __call__(self, k: int, n: int = 20) -> List[ModuleSpec]:
     """
     Args:
       module_specs: list of module_specs to sample from
@@ -161,7 +166,7 @@ def __call__(self,
     # Essentially, split module_specs into k buckets, then define the order of
     # visiting the k buckets such that it approximates the behaviour of having
     # n buckets.
-    specs_len = len(module_specs)
+    specs_len = len(self._module_specs)
     if (specs_len, k, n) not in self._ranges:
       quotient = k // n
       # rev_map maps from bucket # (implicitly via index) to order of visiting.
@@ -177,11 +182,48 @@ def __call__(self,
            math.floor(bucket_size_float * (i + 1))) for i in mapping)
 
     return [
-        module_specs[random.randrange(start, end)]
+        self._module_specs[random.randrange(start, end)]
         for start, end in self._ranges[(specs_len, k, n)]
     ]
 
 
+class CorpusExhaustedError(Exception):
+  pass
+
+
+class SamplerWithoutReplacement(Sampler):
+  """Randomly samples the corpus, without replacement."""
+
+  def __init__(self, module_specs: Tuple[ModuleSpec]):
+    super().__init__(module_specs)
+    self._idx = 0
+    self._shuffle_order()
+
+  def _shuffle_order(self):
+    self._module_specs = tuple(
+        random.sample(self._module_specs, len(self._module_specs)))
+
+  def reset(self):
+    self._shuffle_order()
+    self._idx = 0
+
+  def __call__(self, k: int, n: int = 10) -> List[ModuleSpec]:
+    """
+    Args:
+      k: number of modules to sample
+      n: ignored
+    Raises:
+      CorpusExhaustedError if there are fewer than k elements left to sample in
+      the corpus.
+    """
+    endpoint = self._idx + k
+    if endpoint > len(self._module_specs):
+      raise CorpusExhaustedError()
+    results = self._module_specs[self._idx:endpoint]
+    self._idx = self._idx + k
+    return list(results)
+
+
 class Corpus:
   """Represents a corpus.
 
@@ -230,7 +272,7 @@ def __init__(self,
                additional_flags: Tuple[str, ...] = (),
                delete_flags: Tuple[str, ...] = (),
                replace_flags: Optional[Dict[str, str]] = None,
-               sampler: Sampler = SamplerBucketRoundRobin()):
+               sampler_type: Type[Sampler] = SamplerBucketRoundRobin):
     """
     Prepares the corpus by pre-loading all the CorpusElements and preparing for
     sampling. Command line origin (.cmd file or override) is decided, and final
@@ -252,7 +294,6 @@ def __init__(self,
         matching it. None to include everything.
     """
     self._base_dir = data_path
-    self._sampler = sampler
     # TODO: (b/233935329) Per-corpus *fdo profile paths can be read into
     # {additional|delete}_flags here
     with tf.io.gfile.GFile(
@@ -337,6 +378,10 @@ def get_cmdline(name: str):
               has_thinlto=has_thinlto), module_paths)
     self._module_specs = tuple(
         sorted(contents, key=lambda m: m.size, reverse=True))
+    self._sampler = sampler_type(self._module_specs)
+
+  def reset(self):
+    self._sampler.reset()
 
   def sample(self, k: int, sort: bool = False) -> List[ModuleSpec]:
     """Samples `k` module_specs, optionally sorting by size descending.
@@ -349,7 +394,7 @@ def sample(self, k: int, sort: bool = False) -> List[ModuleSpec]:
     k = min(len(self._module_specs), k)
     if k < 1:
       raise ValueError('Attempting to sample <1 module specs from corpus.')
-    sampled_specs = self._sampler(self._module_specs, k=k)
+    sampled_specs = self._sampler(k=k)
     if sort:
       sampled_specs.sort(key=lambda m: m.size, reverse=True)
     return sampled_specs

compiler_opt/rl/corpus_test.py

@@ -237,6 +237,36 @@ def test_sample(self):
     self.assertEqual(sample[2].name, 'small')
     self.assertEqual(sample[3].name, 'smol')
 
+  def test_sample_without_replacement(self):
+    cps = corpus.create_corpus_for_testing(
+        location=self.create_tempdir(),
+        elements=[
+            corpus.ModuleSpec(name='smol', size=1),
+            corpus.ModuleSpec(name='middle', size=200),
+            corpus.ModuleSpec(name='largest', size=500),
+            corpus.ModuleSpec(name='small', size=100)
+        ],
+        sampler_type=corpus.SamplerWithoutReplacement)
+    samples = []
+    samples.extend(cps.sample(1, sort=True))
+    self.assertLen(samples, 1)
+    samples.extend(cps.sample(1, sort=True))
+    self.assertLen(samples, 2)
+    # Can't sample 3 from the corpus because there are only 2 elements left
+    with self.assertRaises(corpus.CorpusExhaustedError):
+      samples.extend(cps.sample(3, sort=True))
+    # But, we can sample exactly 2 more
+    self.assertLen(samples, 2)
+    samples.extend(cps.sample(2, sort=True))
+    self.assertLen(samples, 4)
+    with self.assertRaises(corpus.CorpusExhaustedError):
+      samples.extend(cps.sample(1, sort=True))
+    samples.sort(key=lambda m: m.size, reverse=True)
+    self.assertEqual(samples[0].name, 'largest')
+    self.assertEqual(samples[1].name, 'middle')
+    self.assertEqual(samples[2].name, 'small')
+    self.assertEqual(samples[3].name, 'smol')
+
   def test_filter(self):
     cps = corpus.create_corpus_for_testing(
         location=self.create_tempdir(),

compiler_opt/rl/local_data_collector_test.py

@@ -116,16 +116,14 @@ def collect_data(self, *args, **kwargs):
 class DeterministicSampler(corpus.Sampler):
   """A corpus sampler that returns modules in order, and can also be reset."""
 
-  def __init__(self):
+  def __init__(self, module_specs: Tuple[corpus.ModuleSpec]):
+    super().__init__(module_specs)
     self._cur_pos = 0
 
-  def __call__(self,
-               module_specs: Tuple[corpus.ModuleSpec],
-               k: int,
-               n: int = 20) -> List[corpus.ModuleSpec]:
+  def __call__(self, k: int, n: int = 20) -> List[corpus.ModuleSpec]:
     ret = []
     for _ in range(k):
-      ret.append(module_specs[self._cur_pos % len(module_specs)])
+      ret.append(self._module_specs[self._cur_pos % len(self._module_specs)])
       self._cur_pos += 1
     return ret
 
@@ -152,16 +150,15 @@ def _test_iterator_fn(data_list):
 
       return _test_iterator_fn
 
-    sampler = DeterministicSampler()
     with LocalWorkerPoolManager(worker_class=MyRunner, count=4) as lwp:
+      cps = corpus.create_corpus_for_testing(
+          location=self.create_tempdir(),
+          elements=[
+              corpus.ModuleSpec(name=f'dummy{i}', size=i) for i in range(100)
+          ],
+          sampler_type=DeterministicSampler)
       collector = local_data_collector.LocalDataCollector(
-          cps=corpus.create_corpus_for_testing(
-              location=self.create_tempdir(),
-              elements=[
-                  corpus.ModuleSpec(name=f'dummy{i}', size=i)
-                  for i in range(100)
-              ],
-              sampler=sampler),
+          cps=cps,
           num_modules=9,
           worker_pool=lwp,
           parser=create_test_iterator_fn(),
@@ -171,7 +168,7 @@ def _test_iterator_fn(data_list):
       # reset the sampler, so the next time we collect, we collect the same
       # modules. We do it before the collect_data call, because that's when
       # we'll re-sample to prefetch the next batch.
-      sampler.reset()
+      cps.reset()
 
       data_iterator, monitor_dict = collector.collect_data(
           policy=_mock_policy, model_id=0)